diff --git a/.circleci/config.yml b/.circleci/config.yml index 6e368a3debe..32d2cf0390c 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1440,6 +1440,7 @@ jobs: TEST_FILES=$(printf "%s\n" \ tests/local_testing/test_dual_cache.py \ tests/local_testing/test_redis_batch_optimizations.py \ + tests/local_testing/test_redis_increment_with_floor.py \ tests/local_testing/test_router_utils.py) echo "$TEST_FILES" | circleci tests run \ --verbose \ @@ -2648,6 +2649,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 +2792,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/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/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-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 4f56e78ddee..c6901411167 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -117,6 +117,9 @@ jobs: - run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl + - name: Run pytest tests/test_litellm_rust with the compiled extension + run: make test-rust-extension + - run: >- uv build --wheel --out-dir panic-dist --config-setting "maturin.build-args=--features panic-test,extension-module" 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 2bc39332817..41678432989 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -29,7 +29,7 @@ Never test structure of code only function of it End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md` -When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions +When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD` When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule @@ -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 -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 +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 the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. 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 @@ -70,7 +70,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests, Always pull before starting any work. The checkout or worktree may be sitting on a stale branch -If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names +If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9ef1d5ae2b8..0443f1bed75 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -315,10 +315,12 @@ Ensure the UI builds successfully before submitting your PR: npm run build ``` +Local lint and budget checks follow origin's current default branch. They refresh it from the remote instead of trusting cached `origin/HEAD`. For an intentional comparison against another branch or commit, use `make check BASE_REF=` or the standalone gate's `--base ` option. An explicit ref can also be used offline once it has been fetched locally. Without an override, unavailable remote metadata stops the check + ## Submitting Your PR 1. **Push your branch**: `git push origin your-feature-branch` -2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`. +2. **Create a PR**: Go to GitHub and open a pull request against the repository's current default branch. Run `python3 scripts/default_branch.py --branch` to check its name 3. **Fill out the PR template**: Provide clear description of changes 4. **Wait for review**: Maintainers will review and provide feedback 5. **Address feedback**: Make requested changes and push updates diff --git a/Makefile b/Makefile index e17fdba3c85..91835e19e3c 100644 --- a/Makefile +++ b/Makefile @@ -4,6 +4,7 @@ .PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \ test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \ test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \ + test-rust-extension \ info lint lint-inner lint-dev lint-checks format \ lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \ lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \ @@ -34,7 +35,7 @@ help: @echo " make lint-basedpyright-budget-update - Ratchet basedpyright limits down by what this branch fixed" @echo " make lint-format - Check ruff format formatting (matches CI)" @echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit" - @echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)" + @echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches the default branch, simulates the merge)" @echo " make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed" @echo " make lint-test-quality - Gate the test suite against test-quality-budget.json" @echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)" @@ -54,12 +55,16 @@ help: @echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)" @echo " make test-integration - Run integration tests" @echo " make test-unit-helm - Run helm unit tests" + @echo " make test-rust-extension - Build the Rust extension and run its public Python tests" @echo "" @echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide" @echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine." UV := uv UV_RUN := $(UV) run --no-sync +BASE_REF ?= +export BASE_REF +RESOLVE_BASE = python3 scripts/default_branch.py --base "$(BASE_REF)" # Machine-wide slot queue for the heavy targets below; python3 + stdlib only, so # it runs before any venv exists. See scripts/gate_slot_lock.py. @@ -67,7 +72,7 @@ GATE_SLOT_LOCK := python3 scripts/gate_slot_lock.py LINT_DEP_INSTALL ?= install-dev LINT_E2E_DEP_INSTALL ?= lint-install -LINT_DEP_BASE ?= lint-fetch-base +LINT_DEP_BASE ?= LINT_JOBS := $(shell sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4) LINT_OUTPUT_SYNC := $(if $(filter output-sync,$(.FEATURES)),--output-sync=target,) @@ -130,10 +135,8 @@ format: install-dev format-check: install-dev cd litellm && $(UV_RUN) ruff format --check --exclude '/enterprise/' . && cd .. -# Single fetch of the PR base so the delta-based gates below share one network round -# trip instead of each re-fetching when chained from `lint`. lint-fetch-base: - git fetch origin litellm_internal_staging + @$(RESOLVE_BASE) # Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated # Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The @@ -150,7 +153,9 @@ lint-install: # recursively, so 'litellm/*.py' covers nested modules and the top-level files that # CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step. lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - @files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \ + @base_ref=$$($(RESOLVE_BASE)) && \ + changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \ + files=$$(printf '%s\n' "$$changed" | grep -v '^litellm/enterprise/' || true) || exit $$?; \ if [ -z "$$files" ]; then \ echo "No changed litellm Python files to format-check."; \ else \ @@ -167,7 +172,9 @@ lint-ruff: $(LINT_DEP_INSTALL) # https://github.com/astral-sh/ruff/discussions/10977 # https://github.com/astral-sh/ruff/discussions/4049 lint-format-changed: install-dev - @git diff origin/main --unified=0 --no-color -- '*.py' | \ + @base_ref=$$($(RESOLVE_BASE)) && \ + diff=$$(git diff "$$base_ref" --unified=0 --no-color -- '*.py') && \ + printf '%s\n' "$$diff" | \ perl -ne '\ if (/^diff --git a\/(.*) b\//) { $$file = $$1; } \ if (/^@@ .* \+(\d+)(?:,(\d+))? @@/) { \ @@ -182,20 +189,22 @@ lint-format-changed: install-dev done lint-ruff-dev: install-dev - @tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \ + @base_ref=$$($(RESOLVE_BASE)) || exit $$?; \ + tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \ cd litellm && \ ($(UV_RUN) ruff check . --output-format=pylint || true) > "$$tmpfile" && \ - $(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch=origin/main && \ + $(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch="$$base_ref" && \ cd .. ; \ rm -f "$$tmpfile" lint-ruff-FULL-dev: install-dev - @files=$$(git diff --name-only origin/main -- '*.py'); \ + @base_ref=$$($(RESOLVE_BASE)) && \ + files=$$(git diff --name-only "$$base_ref" -- '*.py') || exit $$?; \ if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \ else echo "No changed .py files to check."; fi lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/type_check_gate.py --base "$(BASE_REF)" lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL) $(UV_RUN) basedpyright tests/e2e @@ -203,37 +212,37 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL) # Type-discipline budget (mutable collections / casts / type guards / kwargs / # unexplained suppressions), the test-linting.yml step `make lint` used to omit. lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - $(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/type_discipline_gate.py --base "$(BASE_REF)" # Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes, # litellm module-global mutation, credential-gated skips, conftest snapshot # inventory), counted across tests/ the same delta-vs-base way. lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - $(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/test_quality_gate.py --base "$(BASE_REF)" # --update lowers each limit by what this branch fixed since its branch point, so # it needs the base ref fetched to resolve the merge-base. -lint-basedpyright-budget-update: install-dev lint-fetch-base - $(UV_RUN) python scripts/type_check_gate.py --update +lint-basedpyright-budget-update: install-dev + $(UV_RUN) python scripts/type_check_gate.py --update --base "$(BASE_REF)" lint-format: format-check lint-ruff-budget: install-dev - $(UV_RUN) python scripts/ruff_strict_gate.py + $(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)" # Strict gate, invoked the same way CI does in test-linting.yml so a local pass # means the CI check will pass too. lint-gate: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE) - $(UV_RUN) python scripts/ruff_strict_gate.py --base origin/litellm_internal_staging + $(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)" -lint-ruff-budget-update: install-dev lint-fetch-base - $(UV_RUN) python scripts/ruff_strict_gate.py --update +lint-ruff-budget-update: install-dev + $(UV_RUN) python scripts/ruff_strict_gate.py --update --base "$(BASE_REF)" -lint-type-discipline-budget-update: install-dev lint-fetch-base - $(UV_RUN) python scripts/type_discipline_gate.py --update +lint-type-discipline-budget-update: install-dev + $(UV_RUN) python scripts/type_discipline_gate.py --update --base "$(BASE_REF)" -lint-test-quality-budget-update: install-dev lint-fetch-base - $(UV_RUN) python scripts/test_quality_gate.py --update +lint-test-quality-budget-update: install-dev + $(UV_RUN) python scripts/test_quality_gate.py --update --base "$(BASE_REF)" # Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright) lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update @@ -249,14 +258,15 @@ check-import-safety: $(LINT_DEP_INSTALL) # runs the diff-scoped ruff format check, whole-tree ruff check, the strict-rule / # type-discipline / basedpyright budgets as a delta vs the base, then the circular-import # and import-safety checks. Steps that compare against the base resolve it the same way CI -# does (merge-base with origin/litellm_internal_staging). Setup (env sync, Prisma client, +# does (merge-base with origin's current default branch). Setup (env sync, Prisma client, # base fetch) runs once up front; the checks themselves are independent, so a sub-make # fans them out with -j and the fast ones finish under basedpyright's shadow. lint: @$(GATE_SLOT_LOCK) $(MAKE) lint-inner -lint-inner: lint-install lint-fetch-base - $(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks +lint-inner: lint-install + @base_ref=$$($(RESOLVE_BASE)) && \ + $(MAKE) BASE_REF="$$base_ref" -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety @@ -281,6 +291,17 @@ pre-commit: @$(MAKE) check # Testing targets +test-rust-extension: + @temporary=$$(mktemp -d) && \ + trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \ + $(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \ + set -- "$$temporary"/wheels/*.whl && \ + [ "$$#" -eq 1 ] && \ + UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \ + $(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \ + LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \ + "$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust + test: install-test-deps $(UV_RUN) pytest tests/ 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/ci_cd/run_migration.py b/ci_cd/run_migration.py index feec4046ee1..c737050f3a8 100644 --- a/ci_cd/run_migration.py +++ b/ci_cd/run_migration.py @@ -6,12 +6,9 @@ import subprocess import sys from datetime import datetime from pathlib import Path - -import testing.postgresql - +from typing import Final DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE) -DEFAULT_BASE_BRANCH = "litellm_internal_staging" def _find_destructive_statements(sql: str) -> list: @@ -94,31 +91,57 @@ def _print_stale_branch_refusal(base_branch: str, behind: int) -> None: print(banner, file=out) -def _check_branch_freshness(root_dir: Path, base_branch: str) -> None: +def _default_base_branch(root_dir: Path) -> str: + try: + result: Final = subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve().parents[1] / "scripts" / "default_branch.py"), + "--repo-root", + str(root_dir), + "--branch", + ], + check=True, + capture_output=True, + text=True, + timeout=90, + ) + except (OSError, subprocess.SubprocessError) as exc: + _print_freshness_failure( + "default branch", + "Could not discover origin's default branch. Pass --base-branch to choose one.", + exc.stderr if isinstance(exc, subprocess.CalledProcessError) else str(exc), + ) + sys.exit(3) + return result.stdout.strip() + + +def _check_branch_freshness(root_dir: Path, base_branch: str | None = None) -> None: """Fetch origin/ and exit 3 if HEAD is behind it.""" + resolved_branch: Final = base_branch or _default_base_branch(root_dir) cwd = str(root_dir) try: subprocess.run( - ["git", "fetch", "origin", base_branch], + ["git", "fetch", "origin", f"+refs/heads/{resolved_branch}:refs/remotes/origin/{resolved_branch}"], check=True, capture_output=True, text=True, cwd=cwd, ) except FileNotFoundError: - _print_freshness_failure(base_branch, "git executable not found on PATH") + _print_freshness_failure(resolved_branch, "git executable not found on PATH") sys.exit(3) except subprocess.CalledProcessError as e: _print_freshness_failure( - base_branch, - f"`git fetch origin {base_branch}` failed", + resolved_branch, + f"`git fetch origin {resolved_branch}` failed", e.stderr or "", ) sys.exit(3) try: result = subprocess.run( - ["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"], + ["git", "rev-list", "--count", f"HEAD..origin/{resolved_branch}"], check=True, capture_output=True, text=True, @@ -127,23 +150,23 @@ def _check_branch_freshness(root_dir: Path, base_branch: str) -> None: behind = int(result.stdout.strip()) except subprocess.CalledProcessError as e: _print_freshness_failure( - base_branch, - f"`git rev-list HEAD..origin/{base_branch}` failed", + resolved_branch, + f"`git rev-list HEAD..origin/{resolved_branch}` failed", e.stderr or "", ) sys.exit(3) except ValueError: _print_freshness_failure( - base_branch, + resolved_branch, "could not parse commit count from `git rev-list`", ) sys.exit(3) if behind > 0: - _print_stale_branch_refusal(base_branch, behind) + _print_stale_branch_refusal(resolved_branch, behind) sys.exit(3) - print(f"Branch freshness OK: up to date with origin/{base_branch}.") + print(f"Branch freshness OK: up to date with origin/{resolved_branch}.") def _print_destructive_refusal(destructive_lines: list) -> None: @@ -198,7 +221,7 @@ def _print_destructive_refusal(destructive_lines: list) -> None: def create_migration( migration_name: str = None, allow_destructive: bool = False, - base_branch: str = DEFAULT_BASE_BRANCH, + base_branch: str | None = None, skip_freshness_check: bool = False, ): """ @@ -211,7 +234,7 @@ def create_migration( DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this flag, the script exits non-zero and prints guidance. base_branch (str): Branch to check freshness against - (default: "litellm_internal_staging"). + (default: origin's current default branch). skip_freshness_check (bool): Skip the "branch is up to date" check. Only for intentional migrations against an older base. """ @@ -225,6 +248,8 @@ def create_migration( else: _check_branch_freshness(root_dir, base_branch) + import testing.postgresql + try: migrations_dir = ( root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations" @@ -342,9 +367,8 @@ if __name__ == "__main__": ) parser.add_argument( "--base-branch", - default=DEFAULT_BASE_BRANCH, help=( - f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). " + "Branch to check freshness against (default: origin's current default branch). " "The script fetches origin/ and refuses to run if HEAD " "is behind it." ), diff --git a/docker/component_entrypoint.sh b/docker/component_entrypoint.sh index 413957b9929..173afafe1ad 100755 --- a/docker/component_entrypoint.sh +++ b/docker/component_entrypoint.sh @@ -1,5 +1,11 @@ #!/bin/sh +# stale samples from a previous container incarnation would be summed into the aggregate +if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then + mkdir -p "$PROMETHEUS_MULTIPROC_DIR" + rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db +fi + case "$USE_DDTRACE" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED="False" diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index 52ffd117535..f7c918a6827 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -152,6 +152,13 @@ spec: {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} + {{- if .Values.metricsServer.enabled }} + {{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }} + {{- fail "metricsServer.port must differ from service.port" }} + {{- end }} + - name: PROMETHEUS_METRICS_PORT + value: {{ .Values.metricsServer.port | quote }} + {{- end }} {{- if .Values.migrationJob.enabled }} # Schema updates are owned by the dedicated migrations Job; skip # the proxy's startup `prisma db push` so N replicas don't race @@ -189,6 +196,11 @@ spec: - name: http containerPort: {{ .Values.service.port }} protocol: TCP + {{- if .Values.metricsServer.enabled }} + - name: metrics + containerPort: {{ .Values.metricsServer.port }} + protocol: TCP + {{- end }} livenessProbe: httpGet: path: {{ .Values.livenessProbe.path | quote }} diff --git a/helm/litellm-helm/templates/service-metrics.yaml b/helm/litellm-helm/templates/service-metrics.yaml new file mode 100644 index 00000000000..1d23fe39606 --- /dev/null +++ b/helm/litellm-helm/templates/service-metrics.yaml @@ -0,0 +1,17 @@ +{{- if .Values.metricsServer.enabled }} +apiVersion: v1 +kind: Service +metadata: + name: {{ include "litellm.fullname" . }}-metrics + labels: + {{- include "litellm.labels" . | nindent 4 }} +spec: + type: ClusterIP + ports: + - port: {{ .Values.metricsServer.port }} + targetPort: metrics + protocol: TCP + name: metrics + selector: + {{- include "litellm.selectorLabels" . | nindent 4 }} +{{- end }} diff --git a/helm/litellm-helm/templates/servicemonitor.yaml b/helm/litellm-helm/templates/servicemonitor.yaml index 743098deb3f..68083d0da61 100644 --- a/helm/litellm-helm/templates/servicemonitor.yaml +++ b/helm/litellm-helm/templates/servicemonitor.yaml @@ -26,7 +26,7 @@ spec: {{- toYaml .namespaceSelector.matchNames | nindent 4 }} {{- end }} endpoints: - - port: http + - port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }} path: /metrics/ interval: {{ .interval }} scrapeTimeout: {{ .scrapeTimeout }} diff --git a/helm/litellm-helm/tests/metrics_server_tests.yaml b/helm/litellm-helm/tests/metrics_server_tests.yaml new file mode 100644 index 00000000000..085d69ac640 --- /dev/null +++ b/helm/litellm-helm/tests/metrics_server_tests.yaml @@ -0,0 +1,106 @@ +suite: separate metrics server +templates: + - configmap-litellm.yaml + - deployment.yaml + - service.yaml + - service-metrics.yaml + - servicemonitor.yaml +tests: + - it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default + asserts: + - notContains: + path: spec.template.spec.containers[0].ports + content: + name: metrics + any: true + template: deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_METRICS_PORT + any: true + template: deployment.yaml + - lengthEqual: + path: spec.ports + count: 1 + template: service.yaml + - hasDocuments: + count: 0 + template: service-metrics.yaml + + - it: should scrape the proxy port when the metrics server is disabled + template: servicemonitor.yaml + set: + serviceMonitor.enabled: true + asserts: + - equal: + path: spec.endpoints[0].port + value: http + + - it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor + set: + metricsServer.enabled: true + metricsServer.port: 4101 + serviceMonitor.enabled: true + service.type: LoadBalancer + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_METRICS_PORT + value: "4101" + template: deployment.yaml + - contains: + path: spec.template.spec.containers[0].ports + content: + name: metrics + containerPort: 4101 + protocol: TCP + template: deployment.yaml + - lengthEqual: + path: spec.ports + count: 1 + template: service.yaml + - equal: + path: spec.type + value: LoadBalancer + template: service.yaml + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-metrics + template: service-metrics.yaml + - equal: + path: spec.type + value: ClusterIP + template: service-metrics.yaml + - equal: + path: spec.ports + value: + - port: 4101 + targetPort: metrics + protocol: TCP + name: metrics + template: service-metrics.yaml + - equal: + path: spec.selector + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + template: service-metrics.yaml + - equal: + path: spec.endpoints[0].port + value: metrics + template: servicemonitor.yaml + - equal: + path: spec.endpoints[0].path + value: /metrics/ + template: servicemonitor.yaml + + - it: should reject a metrics port equal to the proxy port + template: deployment.yaml + set: + metricsServer.enabled: true + metricsServer.port: 4000 + asserts: + - failedTemplate: + errorMessage: metricsServer.port must differ from service.port diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 637be2322e3..8dc7b967e11 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -180,6 +180,16 @@ proxy_config: general_settings: master_key: os.environ/PROXY_MASTER_KEY +# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT) +# so a scrape never runs on an inference worker. Adds a `metrics` port to the +# container and a dedicated ClusterIP `-metrics` Service, and the +# ServiceMonitor scrapes it instead of the proxy port. The separate port has +# no virtual-key auth: keep it off public ingress. Needs the proxy image +# v1.101.0 or newer. +metricsServer: + enabled: false + port: 4001 + resources: {} # Unset by default so the chart installs on small clusters such as Minikube, and so an diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index be6b9093f53..c459512c7b9 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -441,3 +441,5 @@ ImplementationSpecific {{- .pathType -}} {{- end -}} {{- end -}} + +{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}} diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 5030ba2c9dc..9cb6b07e77b 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -64,14 +64,25 @@ spec: {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsEnv" . | nindent 12 }} {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + {{- if eq (int .Values.gateway.metricsServer.port) 4000 }} + {{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }} + {{- end }} + - name: PROMETHEUS_MULTIPROC_DIR + value: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + {{- end }} {{- include "litellm.envFrom" .Values.gateway | nindent 10 }} - {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }} volumeMounts: {{- if .Values.gateway.config.create }} - name: gateway-config mountPath: /app/config/config.yaml subPath: config.yaml {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: prometheus-multiproc + mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + {{- end }} {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }} {{- end }} @@ -97,16 +108,54 @@ spec: {{- end }} resources: {{- toYaml .Values.gateway.resources | nindent 12 }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: metrics + image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.gateway.image.pullPolicy }} + {{- with .Values.gateway.securityContext }} + securityContext: + {{- toYaml . | nindent 12 }} + {{- end }} + command: + - python + - -m + - litellm.proxy.prometheus_metrics_server + - --port + - {{ .Values.gateway.metricsServer.port | quote }} + env: + - name: PROMETHEUS_MULTIPROC_DIR + value: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + ports: + - name: metrics + containerPort: {{ .Values.gateway.metricsServer.port }} + protocol: TCP + volumeMounts: + - name: prometheus-multiproc + mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }} + readinessProbe: + tcpSocket: { port: metrics } + periodSeconds: 10 + livenessProbe: + tcpSocket: { port: metrics } + periodSeconds: 15 + failureThreshold: 6 + resources: + {{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }} + {{- end }} {{- with .Values.gateway.extraContainers }} {{- tpl (toYaml .) $ | nindent 8 }} {{- end }} - {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }} + {{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }} volumes: {{- if .Values.gateway.config.create }} - name: gateway-config configMap: name: {{ include "litellm.gateway.fullname" . }}-config {{- end }} + {{- if .Values.gateway.metricsServer.enabled }} + - name: prometheus-multiproc + emptyDir: {} + {{- end }} {{- if .Values.billingMetrics.enabled }} {{- include "litellm.billingMetricsVolumes" . | nindent 8 }} {{- end }} diff --git a/helm/litellm/templates/gateway/service-metrics.yaml b/helm/litellm/templates/gateway/service-metrics.yaml new file mode 100644 index 00000000000..ad9bc05a9fd --- /dev/null +++ b/helm/litellm/templates/gateway/service-metrics.yaml @@ -0,0 +1,18 @@ +{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }} +apiVersion: v1 +kind: Service +metadata: + name: {{ include "litellm.gateway.fullname" . }}-metrics + labels: + {{- include "litellm.commonLabels" . | nindent 4 }} + app.kubernetes.io/component: gateway +spec: + type: ClusterIP + ports: + - port: {{ .Values.gateway.metricsServer.port }} + targetPort: metrics + protocol: TCP + name: metrics + selector: + {{- include "litellm.gateway.selectorLabels" . | nindent 4 }} +{{- end }} diff --git a/helm/litellm/tests/metrics_server_tests.yaml b/helm/litellm/tests/metrics_server_tests.yaml new file mode 100644 index 00000000000..0e7d9d9e9ee --- /dev/null +++ b/helm/litellm/tests/metrics_server_tests.yaml @@ -0,0 +1,148 @@ +suite: test gateway metrics sidecar +templates: + - gateway/configmap.yaml + - gateway/deployment.yaml + - gateway/service.yaml + - gateway/service-metrics.yaml +values: + - ./values/required.yaml +tests: + - it: adds no sidecar, volume, env or service port when the metrics server is off + asserts: + - lengthEqual: + path: spec.template.spec.containers + count: 1 + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_MULTIPROC_DIR + any: true + template: gateway/deployment.yaml + - notContains: + path: spec.template.spec.volumes + content: + name: prometheus-multiproc + any: true + template: gateway/deployment.yaml + - lengthEqual: + path: spec.ports + count: 1 + template: gateway/service.yaml + - hasDocuments: + count: 0 + template: gateway/service-metrics.yaml + + - it: runs the metrics server as a sidecar over a shared multiproc dir and exposes it on a ClusterIP metrics service + set: + gateway.metricsServer.enabled: true + gateway.metricsServer.port: 4101 + gateway.service.type: LoadBalancer + gateway.image.tag: v1.101.0 + asserts: + - contains: + path: spec.template.spec.containers[0].env + content: + name: PROMETHEUS_MULTIPROC_DIR + value: /tmp/litellm_prometheus_multiproc + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[0].volumeMounts + content: + name: prometheus-multiproc + mountPath: /tmp/litellm_prometheus_multiproc + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].name + value: metrics + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].image + value: ghcr.io/berriai/litellm-gateway:v1.101.0 + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].command + value: + - python + - -m + - litellm.proxy.prometheus_metrics_server + - --port + - "4101" + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].env + value: + - name: PROMETHEUS_MULTIPROC_DIR + value: /tmp/litellm_prometheus_multiproc + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].ports + value: + - name: metrics + containerPort: 4101 + protocol: TCP + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].volumeMounts + value: + - name: prometheus-multiproc + mountPath: /tmp/litellm_prometheus_multiproc + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].readinessProbe.tcpSocket.port + value: metrics + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].livenessProbe.tcpSocket.port + value: metrics + template: gateway/deployment.yaml + - equal: + path: spec.template.spec.containers[1].resources.requests.cpu + value: 50m + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.volumes + content: + name: prometheus-multiproc + emptyDir: {} + template: gateway/deployment.yaml + - lengthEqual: + path: spec.ports + count: 1 + template: gateway/service.yaml + - equal: + path: spec.type + value: LoadBalancer + template: gateway/service.yaml + - equal: + path: metadata.name + value: RELEASE-NAME-litellm-gateway-metrics + template: gateway/service-metrics.yaml + - equal: + path: spec.type + value: ClusterIP + template: gateway/service-metrics.yaml + - equal: + path: spec.ports + value: + - port: 4101 + targetPort: metrics + protocol: TCP + name: metrics + template: gateway/service-metrics.yaml + - equal: + path: spec.selector + value: + app.kubernetes.io/name: litellm + app.kubernetes.io/instance: RELEASE-NAME + app.kubernetes.io/component: gateway + template: gateway/service-metrics.yaml + + - it: rejects a metrics port equal to the gateway port + template: gateway/deployment.yaml + set: + gateway.metricsServer.enabled: true + gateway.metricsServer.port: 4000 + asserts: + - failedTemplate: + errorMessage: gateway.metricsServer.port must differ from the gateway port 4000 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 6c9fb9440c7..b5d535c992d 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -268,6 +268,22 @@ gateway: config: create: true proxy_config: {} + # Serve Prometheus /metrics from a `metrics` sidecar container (same image, + # `python -m litellm.proxy.prometheus_metrics_server`) that aggregates the + # workers' PROMETHEUS_MULTIPROC_DIR samples over a shared emptyDir, so a + # scrape never runs on an inference worker. Adds a `metrics` port to the pod + # and a dedicated ClusterIP `-metrics` Service; point your scrape + # config at it. The port has no virtual-key auth: keep it off public ingress. + # Needs the gateway image v1.101.0 or newer. + metricsServer: + enabled: false + port: 4001 + resources: + requests: + cpu: 50m + memory: 128Mi + limits: + memory: 512Mi image: repository: ghcr.io/berriai/litellm-gateway tag: "" # defaults to .Chart.AppVersion diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql new file mode 100644 index 00000000000..5503167ce09 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql @@ -0,0 +1,3 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" +ADD COLUMN IF NOT EXISTS "classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0, +ADD COLUMN IF NOT EXISTS "classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 06ac177cca4..ccbab0fef10 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1509,6 +1509,8 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 71e7e9c683b..2145f891318 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -8,14 +8,10 @@ import tempfile import time from dataclasses import dataclass, replace from pathlib import Path -from typing import Optional +from typing import TYPE_CHECKING, Final, Optional from litellm_proxy_extras import prisma_toolchain from litellm_proxy_extras._logging import logger -from litellm_proxy_extras.replica_identity import ( - REPLICA_IDENTITY_FULL_ENV_VAR, - apply_replica_identity_full, -) from litellm_proxy_extras.prisma_toolchain import ( PRISMA_COMMAND_TIMEOUT_ENV_VAR, PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR, @@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import ( prisma_command_timeout, prisma_migrate_deploy_timeout, ) +from litellm_proxy_extras.replica_identity import ( + REPLICA_IDENTITY_FULL_ENV_VAR, + apply_replica_identity_full, +) + +if TYPE_CHECKING: + import psycopg + import psycopg.sql def str_to_bool(value: Optional[str]) -> bool: @@ -46,6 +50,28 @@ def _get_prisma_env() -> dict: _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") _MIGRATION_DEADLOCK_MARKER = "deadlock detected" +INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") +_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") +_INVALID_LITELLM_INDEXES_SQL: Final = ( + "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " + "FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_class t ON t.oid = i.indrelid " + "JOIN pg_namespace n ON n.oid = t.relnamespace " + "WHERE NOT i.indisvalid " + " AND c.relkind = 'i' " + " AND n.nspname = %s " + " AND t.relname LIKE %s " + " AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) " + "ORDER BY c.relname" +) + + +@dataclass(frozen=True, slots=True) +class _InvalidIndex: + schema: str + name: str + table_size: str MAX_MIGRATE_DEPLOY_ATTEMPTS = 4 @@ -624,7 +650,7 @@ class ProxyExtrasDBManager: def _strip_prisma_query_params(url: str) -> str: """Remove Prisma-specific query params (connection_limit, pool_timeout, schema, etc.) from DATABASE_URL so psycopg can parse it.""" - from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode + from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse parsed = urlparse(url) if not parsed.query: @@ -645,7 +671,7 @@ class ProxyExtrasDBManager: "target_session_attrs", } kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params] - return urlunparse(parsed._replace(query=urlencode(kept))) + return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote))) @staticmethod def _warn_if_db_ahead_of_head(migrations_dir: str) -> None: @@ -719,6 +745,95 @@ class ProxyExtrasDBManager: ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""), ) + @staticmethod + def _invalid_litellm_indexes( + conn: "psycopg.Connection[tuple[str, str, str]]", schema: str + ) -> tuple[_InvalidIndex, ...]: + rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall() + return tuple(_InvalidIndex(*row) for row in rows) + + @staticmethod + def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]: + from psycopg import sql + + target: Final = sql.Identifier(index.schema, index.name) + if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name): + return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover" + return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt" + + @staticmethod + def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None: + import psycopg + + statement, action = ProxyExtrasDBManager._index_repair(index) + try: + conn.execute(statement) + except psycopg.Error as e: + logger.warning( + "Could not repair invalid index %s.%s, will retry on the next startup. " + "If this keeps happening, run `%s` by hand as the index owner. Error: %s", + index.schema, + index.name, + statement.as_string(conn), + e, + ) + return + logger.info("%s invalid index %s.%s", action, index.schema, index.name) + + @staticmethod + def repair_invalid_indexes(lock_timeout: str = "30s") -> bool: + """Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left + INVALID (a migration deadlock between replicas is the usual cause; the + retried migration skips them because of IF NOT EXISTS). Never raises: + returns True when no invalid index remains, False when the repair was + skipped or failed and will be retried on the next startup. Looks in the + schema DATABASE_URL names, the only URL Prisma migrates through, but + connects over DIRECT_URL when set: the session settings, the advisory + lock and REINDEX CONCURRENTLY all need one server session, which a + transaction pooler does not give.""" + prisma_url: Final = os.getenv("DATABASE_URL") + if not prisma_url: + return False + + try: + import psycopg + from psycopg import sql + except ImportError: + logger.warning( + "psycopg is not installed; skipping the invalid index check. " + "Install the litellm[extra_proxy] extra, which includes psycopg." + ) + return False + + schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url) + try: + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: + conn.execute("SET statement_timeout = 0") + conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout))) + found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + if not found: + return True + logger.warning( + "Found %d invalid index(es) left by an interrupted CREATE INDEX " + "CONCURRENTLY, rebuilding: %s", + len(found), + ", ".join(f"{index.name} (table size {index.table_size})" for index in found), + ) + lock_row: Final = conn.execute( + "SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,) + ).fetchone() + if lock_row is None or not lock_row[0]: + logger.info("Another replica is already rebuilding the invalid indexes, skipping") + return False + for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema): + ProxyExtrasDBManager._repair_index(conn, index) + remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + except psycopg.Error as e: + logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e) + return False + return not remaining + @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: """ @@ -994,6 +1109,7 @@ class ProxyExtrasDBManager: use_migrate=use_migrate, use_v2_resolver=use_v2_resolver ) if migrated: + ProxyExtrasDBManager.repair_invalid_indexes() ProxyExtrasDBManager.apply_replica_identity_full_if_requested() return migrated diff --git a/litellm-proxy-extras/migration_runbook.md b/litellm-proxy-extras/migration_runbook.md index a277441b164..b1e9236e520 100644 --- a/litellm-proxy-extras/migration_runbook.md +++ b/litellm-proxy-extras/migration_runbook.md @@ -48,7 +48,7 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n ## What It Does -1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check)) +1. **Verifies the current branch is up to date with origin's current default branch** (see [Branch freshness](#branch-freshness-check)) 2. Creates temp PostgreSQL DB 3. Applies existing migrations 4. Compares with `schema.prisma` @@ -57,11 +57,11 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n ## Branch Freshness Check -Before generating anything, `run_migration.py` runs `git fetch origin ` and refuses to proceed if `HEAD` is behind `origin/`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense. +Before generating anything, `run_migration.py` runs `git fetch origin ` and refuses to proceed if `HEAD` is behind `origin/`. The default base is discovered from origin's advertised HEAD on each run, so an existing clone follows a default-branch change without trusting cached `origin/HEAD`. If discovery or fetching fails, migration generation stops. A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense. Flags: -- `--base-branch ` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`. +- `--base-branch ` — check against a different base (e.g. a release branch). Defaults to origin's current default branch - `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base. When the guard fires: @@ -69,8 +69,9 @@ When the guard fires: 1. Update your branch: ```bash - git fetch origin && git rebase origin/litellm_internal_staging - # or git merge origin/litellm_internal_staging — whichever matches your workflow + base_branch=$(python3 scripts/default_branch.py --branch) && + git fetch origin "+refs/heads/$base_branch:refs/remotes/origin/$base_branch" && + git rebase "origin/$base_branch" ``` 2. Re-run `run_migration.py`. diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 82d31fec373..91b4e4a7ba1 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.94" +version = "0.4.95" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.94" +version = "0.4.95" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..62477dd6264 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) @@ -2001,6 +2002,9 @@ if TYPE_CHECKING: from .llms.hosted_vllm.responses.transformation import ( HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig, ) + from .llms.fireworks_ai.responses.transformation import ( + FireworksAIResponsesAPIConfig as FireworksAIResponsesAPIConfig, + ) from .llms.github_copilot.chat.transformation import ( GithubCopilotConfig as GithubCopilotConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index e9199e1ec80..dc323c8cc15 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -237,6 +237,7 @@ LLM_CONFIG_NAMES: Final = ( "XAIResponsesAPIConfig", "LiteLLMProxyResponsesAPIConfig", "HostedVLLMResponsesAPIConfig", + "FireworksAIResponsesAPIConfig", "VolcEngineResponsesAPIConfig", "PerplexityResponsesConfig", "DatabricksResponsesAPIConfig", @@ -957,6 +958,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.hosted_vllm.responses.transformation", "HostedVLLMResponsesAPIConfig", ), + "FireworksAIResponsesAPIConfig": ( + ".llms.fireworks_ai.responses.transformation", + "FireworksAIResponsesAPIConfig", + ), "VolcEngineResponsesAPIConfig": ( ".llms.volcengine.responses.transformation", "VolcEngineResponsesAPIConfig", diff --git a/litellm/batches/main.py b/litellm/batches/main.py index c8360a81c7a..77a4fdebf16 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -319,6 +319,7 @@ def create_batch( timeout=timeout, max_retries=optional_params.max_retries, create_batch_data=_create_batch_request, + custom_endpoint=optional_params.get("custom_endpoint"), ) else: raise litellm.exceptions.BadRequestError( diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index aaee7188d86..106c1580110 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -20,6 +20,8 @@ from contextvars import ContextVar from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( @@ -80,11 +82,29 @@ class _AsyncRedisCommands(Protocol): def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... + def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( {"", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"} ) +_INCREMENT_WITH_FLOOR_LUA: Final = ( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end " + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " + "return count" +) + +_LUA_COUNT: Final = TypeAdapter(int) +_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) + + +def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: + return _OPTIONAL_COUNTS.validate_python( + tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) + ) + def _get_call_stack_info(num_frames: int = 2) -> str: """ @@ -736,6 +756,43 @@ class RedisCache(BaseCache): ) raise e + @_redis_circuit_breaker_guard_sync + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call. + + A counter whose key expired while a request was still in flight would otherwise be + recreated negative by that request's decrement. Clamping inside the same call is what + keeps it safe: a separate corrective write could land after another pod's increment and + erase it. + + The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was + created rather than ``ttl`` after it was last touched. Refreshing it on every touch + would keep a count a dead worker never decremented alive for as long as the group + takes traffic. Returns the resulting count. + """ + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval + _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl + ) + return _LUA_COUNT.validate_python(count) + + @_redis_circuit_breaker_guard_sync + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Read integer counters for ``key_list``, in order, raising when Redis cannot answer. + + ``batch_get_cache`` swallows every failure and returns an empty dict, which the caller + cannot tell apart from "every counter is unset". A caller that has to fall back to its + own numbers when Redis is unreachable needs the failure, not a dict of zeros. + """ + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys)) + + @_redis_circuit_breaker_guard + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Async twin of ``batch_get_counts``, raising on failure the same way.""" + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() @@ -1241,6 +1298,14 @@ class RedisCache(BaseCache): result = result.decode() return float(result) + @_redis_circuit_breaker_guard + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees.""" + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl) + return _LUA_COUNT.validate_python(count) + async def flush_cache_buffer(self): print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}") await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 87350b5479c..4fe069b0b7d 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -370,7 +370,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): and isinstance(tool_call.get("custom"), dict) ) - for msg in messages: + leading_system_count: Final = next( + (index for index, msg in enumerate(messages) if msg.get("role") != "system"), + len(messages), + ) + + for index, msg in enumerate(messages): role = msg.get("role") content = msg.get("content", "") tool_calls = msg.get("tool_calls") @@ -378,7 +383,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if role == "system": # Extract system message as instructions - if isinstance(content, str): + if isinstance(content, str) and index < leading_system_count: if instructions: # Concatenate multiple system prompts with a space instructions = f"{instructions} {content}" diff --git a/litellm/constants.py b/litellm/constants.py index d53686e5e5b..8a6a044ff7b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16 MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)) +budget_reservation_disabled_info_emitted = False DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)) DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512)) SQS_SEND_MESSAGE_ACTION: Final = "SendMessage" @@ -72,6 +73,9 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096)) DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3)) DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1)) DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5)) +DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float( + os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1") +) DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) @@ -1457,6 +1461,8 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata" OLD_LITELLM_METADATA_FIELD: Final = "metadata" RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name" SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl" +OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens", "max_output_tokens"}) +CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling" CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags" INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9a9d2ceda03..c9c7df4d75e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2414,6 +2414,46 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): ) +_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"}) + + +class _ResponsesWsEventResponse(BaseModel): + usage: Mapping[str, object] | None = None + + +class _ResponsesWsEvent(BaseModel): + type: str = "" + response: _ResponsesWsEventResponse | None = None + + +class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): + @staticmethod + def collect_usage_from_responses_ws_results( + results: Sequence[Mapping[str, object]], + ) -> tuple[Usage, ...]: + events: Final = tuple(_ResponsesWsEvent.model_validate(result) for result in results) + return tuple( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the realtime processor uses + event.response.usage + ) + for event in events + if event.type in _RESPONSES_WS_BILLABLE_EVENT_TYPES + and event.response is not None + and event.response.usage is not None + ) + + @staticmethod + def collect_and_combine_usage_from_responses_ws_results( + results: Sequence[Mapping[str, object]], + ) -> Usage: + collected_usage_objects: Final = ResponsesWebSocketTokenUsageProcessor.collect_usage_from_responses_ws_results( + results + ) + return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects( + list(collected_usage_objects) # mutable-ok: combine_usage_objects requires a list parameter + ) + + _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 16202321709..f9215267bf3 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -338,6 +338,7 @@ class Timeout(openai.APITimeoutError): num_retries: int | None = None, headers: dict | None = None, exception_status_code: int | None = None, + response: httpx.Response | None = None, ): request: Final = httpx.Request( method="POST", @@ -352,6 +353,8 @@ class Timeout(openai.APITimeoutError): self.max_retries = max_retries self.num_retries = num_retries self.headers = headers + if response is not None: + self.response = response # custom function to convert to str def __str__(self): diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 3519240dda9..4f9b18713d0 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement @@ -23,6 +25,7 @@ from litellm.integrations.prompt_management_base import PromptManagementClient from litellm.litellm_core_utils.prompt_templates.common_utils import ( with_prompt_cache_breakpoint, ) +from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request from litellm.types.integrations.anthropic_cache_control_hook import ( GATEWAY_INJECTED_CACHE_METADATA_KEY, GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT, @@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset( ) OPENAI_API_HOST: Final = "api.openai.com" OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE") +_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues +def _validated_object_mapping(value: object) -> dict[object, object] | None: + try: + return _OBJECT_MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _validated_object_list(value: object) -> list[object] | None: + try: + return _OBJECT_LIST_ADAPTER.validate_python(value) + except ValidationError: + return None + + def supports_openai_prompt_cache_breakpoint(model: str) -> bool: model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model) if model_map_flag is not None: @@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_ class AnthropicCacheControlHook(CustomPromptManagement): + @staticmethod + def _request_value(request_kwargs: object, key: str) -> object: + request_mapping: Final = _validated_object_mapping(request_kwargs) + if request_mapping is None: + return None + return request_mapping.get(key) + + @staticmethod + def _request_user_agent(request_kwargs: object) -> str | None: + proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request") + proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request) + if proxy_server_request_mapping is None: + return None + headers: Final = proxy_server_request_mapping.get("headers") + headers_mapping: Final = _validated_object_mapping(headers) + if headers_mapping is None: + return None + user_agent: Final = next( + (value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"), + None, + ) + return user_agent if isinstance(user_agent, str) else None + + @staticmethod + def _request_system(request_kwargs: object) -> str | list[object] | None: + system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system") + if isinstance(system, str): + return system + return _validated_object_list(system) + def get_chat_completion_prompt( self, model: str, @@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): points: Sequence[CacheControlInjectionPoint], messages: list[AllMessageValues], tools: list[object] | None, + cache_control: object, model: str, custom_llm_provider: str | None, api_base: object, prompt_cache_options: object, ) -> Sequence[Mapping[str, object]] | None: - if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools): + if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control): return None return AnthropicCacheControlHook._stamped_with_dialect( points, model, custom_llm_provider, api_base, prompt_cache_options @@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): messages: list[AllMessageValues], system: str | list | None, tools: list | None, + cache_control: object = None, ) -> bool: """Whether configured injection points must yield to client-set cache_control. @@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement): """ if all(point.get("_litellm_judged") for point in points): return False - return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools) + return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control) @staticmethod def _request_has_cache_control( messages: list[AllMessageValues], system: str | list | None, tools: list | None = None, + cache_control: object = None, ) -> bool: """Return True if the request already carries any client-supplied cache_control. @@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): carry the mark either at the top level (Anthropic shape) or nested under ``function`` (OpenAI shape); the Anthropic chat transform accepts both. """ + if cache_control is not None: + return True if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0: return True if tools is not None: @@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider: str | None, tools: list | None = None, enable_prompt_caching: bool | None = None, + cache_control: object = None, + request_kwargs: object = None, ) -> list[CacheControlInjectionPoint]: """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on. @@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement): if not supports_prompt_caching(model=model, custom_llm_provider=provider): return [] - if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools): + if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control): + return [] + + if is_claude_code_one_shot_subagent_request( + messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs) + ): return [] control: Final = AnthropicCacheControlHook._default_control() @@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): models: Iterable[str], tools: list[AllToolParamValues] | None = None, enable_prompt_caching: bool | None = None, + request_kwargs: object = None, ) -> list[AllMessageValues]: """Return the messages auto prompt caching will send, default breakpoints included. @@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): for candidate in ( AnthropicCacheControlHook.get_default_injection_points( messages=messages, - system=None, model=model, custom_llm_provider=None, tools=tools, enable_prompt_caching=enable_prompt_caching, + system=AnthropicCacheControlHook._request_system(request_kwargs), + cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"), + request_kwargs=request_kwargs, ) for model in models ) @@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): non_default_params["cache_control_injection_points"], messages, tools, + non_default_params.get("cache_control"), model, custom_llm_provider, api_base, @@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): custom_llm_provider=custom_llm_provider, tools=tools, enable_prompt_caching=enable_prompt_caching, + cache_control=non_default_params.get("cache_control"), + request_kwargs=non_default_params, ) if points: non_default_params["cache_control_injection_points"] = points @@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement): enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy bool | None, kwargs.pop("enable_prompt_caching", None) ) + cache_control: Final = kwargs.get("cache_control") configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None) ) - if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools): + if configured and AnthropicCacheControlHook._should_stand_down( + configured, typed_messages, system, tools, cache_control + ): return messages, system injection_points: list[CacheControlInjectionPoint] = configured or [] if not injection_points and model is not None: @@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement): model=model, custom_llm_provider=custom_llm_provider, enable_prompt_caching=enable_prompt_caching, + cache_control=cache_control, + request_kwargs=kwargs, ) if not injection_points: return messages, system diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 7a2295a35ae..b98eef8329b 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -356,6 +356,12 @@ "description": "OpenTelemetry collector endpoint URL", "required": true }, + "otel_traces_endpoint": { + "type": "text", + "ui_name": "Traces Endpoint URL", + "description": "Complete trace export URL used verbatim when the collector does not serve /v1/traces (OTel v2 only)", + "required": false + }, "otel_headers": { "type": "text", "ui_name": "Headers", diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2d66a280663..37d6a7e793d 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger): event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None, supported_event_hooks: list[GuardrailEventHooks], ) -> None: + allowed_hooks: Final = frozenset(supported_event_hooks) | ( + frozenset((GuardrailEventHooks.logging_only,)) + if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks + else frozenset() + ) + def _validate_event_hook_list_is_in_supported_event_hooks( event_hook: list[GuardrailEventHooks] | list[str], supported_event_hooks: list[GuardrailEventHooks], @@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger): for hook in event_hook: if isinstance(hook, str): hook = GuardrailEventHooks(hook) - if hook not in supported_event_hooks: + if hook not in allowed_hooks: raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}") if event_hook is None: @@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger): default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default] _validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks) elif isinstance(event_hook, GuardrailEventHooks): - if event_hook not in supported_event_hooks: + if event_hook not in allowed_hooks: raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}") @staticmethod @@ -773,7 +779,7 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail - def _deployment_pre_call_target(self) -> "CustomLogger": + def _deployment_hook_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self try: @@ -802,7 +808,7 @@ class CustomGuardrail(CustomLogger): # CHECK IF GUARDRAIL REJECTS THE REQUEST if call_type == CallTypes.completion or call_type == CallTypes.acompletion: - target: Final = self._deployment_pre_call_target() + target: Final = self._deployment_hook_target() if target is not self: kwargs["guardrail_to_apply"] = self result: Final = await target.async_pre_call_hook( @@ -845,7 +851,9 @@ class CustomGuardrail(CustomLogger): return None # CHECK IF GUARDRAIL REJECTS THE REQUEST - result: Final = await self.async_post_call_success_hook( + target: Final = self._deployment_hook_target() + hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data + result: Final = await target.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth( user_id=request_data.get("user_api_key_user_id"), team_id=request_data.get("user_api_key_team_id"), @@ -853,7 +861,7 @@ class CustomGuardrail(CustomLogger): api_key=request_data.get("user_api_key_hash"), request_route=request_data.get("user_api_key_request_route"), ), - data=request_data, + data=hook_request_data, response=response, ) diff --git a/litellm/integrations/otel/model/config.py b/litellm/integrations/otel/model/config.py index 9e3064c2bff..e3501d0ee94 100644 --- a/litellm/integrations/otel/model/config.py +++ b/litellm/integrations/otel/model/config.py @@ -72,6 +72,14 @@ class ExporterSpec(BaseModel): description="console | in_memory | otlp_http | otlp_grpc | ", ) endpoint: str | None = None + traces_endpoint: str | None = Field( + default=None, + description=( + "Complete OTLP/HTTP trace URL, used verbatim. Set this when the " + "collector serves traces on a path other than ``/v1/traces``; " + "``endpoint`` is a base URL the signal path is appended to." + ), + ) headers: str | None = None owner: ExporterOwner | None = Field( default=None, @@ -127,6 +135,14 @@ class OpenTelemetryV2Config(BaseSettings): default=None, validation_alias=AliasChoices("OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"), ) + traces_endpoint: str | None = Field( + default=None, + validation_alias=AliasChoices("OTEL_TRACES_ENDPOINT", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"), + description=( + "Complete OTLP/HTTP trace URL for the single-destination shorthand, " + "used verbatim instead of ``endpoint`` + ``/v1/traces``." + ), + ) headers: str | None = Field( default=None, validation_alias=AliasChoices("OTEL_HEADERS", "OTEL_EXPORTER_OTLP_HEADERS"), @@ -250,7 +266,7 @@ class OpenTelemetryV2Config(BaseSettings): @model_validator(mode="after") def _normalize(self) -> "OpenTelemetryV2Config": # An endpoint with the default exporter kind implies OTLP/HTTP. - if self.endpoint and self.exporter == "console": + if (self.endpoint or self.traces_endpoint) and self.exporter == "console": self.exporter = "otlp_http" # When no explicit destinations are given, fold the single-destination # shorthand into one spec so the provider always has a destination. @@ -259,6 +275,7 @@ class OpenTelemetryV2Config(BaseSettings): ExporterSpec( kind=self.exporter, endpoint=self.endpoint, + traces_endpoint=self.traces_endpoint, headers=self.headers, ) ] diff --git a/litellm/integrations/otel/plumbing/providers.py b/litellm/integrations/otel/plumbing/providers.py index fb74ff85e5b..d8da53017ab 100644 --- a/litellm/integrations/otel/plumbing/providers.py +++ b/litellm/integrations/otel/plumbing/providers.py @@ -170,7 +170,7 @@ def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter: ) return HTTPExporter( - endpoint=_otlp_traces_endpoint(spec.endpoint), + endpoint=spec.traces_endpoint or _otlp_traces_endpoint(spec.endpoint), headers=parse_headers(spec.headers), ) if kind in _OTLP_GRPC_KINDS: @@ -201,7 +201,14 @@ def build_span_exporter(config: OpenTelemetryV2Config) -> SpanExporter: ``exporter`` / ``endpoint`` / ``headers`` fields. To configure multiple exporters, populate ``config.exporters`` directly. """ - return _exporter_from_spec(ExporterSpec(kind=config.exporter, endpoint=config.endpoint, headers=config.headers)) + return _exporter_from_spec( + ExporterSpec( + kind=config.exporter, + endpoint=config.endpoint, + traces_endpoint=config.traces_endpoint, + headers=config.headers, + ) + ) def _otlp_metrics_endpoint(endpoint: str | None) -> str | None: 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/credential_accessor.py b/litellm/litellm_core_utils/credential_accessor.py index 7071750b970..7b4e8240b69 100644 --- a/litellm/litellm_core_utils/credential_accessor.py +++ b/litellm/litellm_core_utils/credential_accessor.py @@ -7,16 +7,19 @@ from litellm.types.utils import CredentialItem class CredentialAccessor: + @staticmethod + def find_credential(credential_name: str) -> CredentialItem | None: + return next( + (credential for credential in litellm.credential_list if credential.credential_name == credential_name), + None, + ) + @staticmethod def get_credential_values(credential_name: str) -> dict: """Safe accessor for credentials.""" - if not litellm.credential_list: - return {} - for credential in litellm.credential_list: - if credential.credential_name == credential_name: - return credential.credential_values.copy() - return {} + credential: Final = CredentialAccessor.find_credential(credential_name) + return {} if credential is None else credential.credential_values.copy() @staticmethod def upsert_credentials(credentials: list[CredentialItem]): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 8f8c955d971..82708d412c9 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -860,6 +860,7 @@ def _map_bedrock_exception( message=mantle_context_window_message, model=model, llm_provider=custom_llm_provider, + response=getattr(original_exception, "response", None), ) if ( "too many tokens" in error_str @@ -873,6 +874,7 @@ def _map_bedrock_exception( message=f"BedrockException: Context Window Error - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif "Conversation blocks and tool result blocks cannot be provided in the same turn." in error_str: raise BadRequestError( @@ -924,12 +926,14 @@ def _map_bedrock_exception( message=f"BedrockException: Timeout Error - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif "Could not process image" in error_str: raise litellm.InternalServerError( message=f"BedrockException - {error_str}", model=model, llm_provider="bedrock", + response=getattr(original_exception, "response", None), ) elif hasattr(original_exception, "status_code"): if original_exception.status_code == 500: @@ -937,10 +941,7 @@ def _map_bedrock_exception( message=f"BedrockException - {original_exception.message}", llm_provider="bedrock", model=model, - response=httpx.Response( - status_code=500, - request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), - ), + response=getattr(original_exception, "response", None), ) elif original_exception.status_code == 401: raise AuthenticationError( @@ -969,6 +970,7 @@ def _map_bedrock_exception( model=model, llm_provider=custom_llm_provider, litellm_debug_info=extra_information, + response=getattr(original_exception, "response", None), ) elif original_exception.status_code == 422: raise BadRequestError( @@ -1001,6 +1003,7 @@ def _map_bedrock_exception( llm_provider=custom_llm_provider, litellm_debug_info=extra_information, exception_status_code=original_exception.status_code, + response=getattr(original_exception, "response", None), ) 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 ca2cca5360f..e3c94866985 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -48,6 +48,7 @@ from litellm.constants import ( ) from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, + ResponsesWebSocketTokenUsageProcessor, _select_model_name_for_cost_calc, ) from litellm.exceptions import ( @@ -2028,6 +2029,17 @@ class Logging(LiteLLMLoggingBaseClass): results=result, ) + elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped + combined_ws_usage: Final = ( + ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( + results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream + ) + ) + logging_result = LiteLLMRealtimeStreamLoggingObject( + usage=combined_ws_usage, + results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream + ) + elif ( self.call_type == CallTypes.llm_passthrough_route.value or self.call_type == CallTypes.allm_passthrough_route.value diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 9432fefc368..68dc27ec25e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -514,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. @@ -524,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) """ @@ -551,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 @@ -662,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): @@ -673,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, @@ -1416,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/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py index f1ae6492e4e..d04abcb6e7b 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py @@ -1,7 +1,8 @@ +from collections.abc import Mapping from typing import Final -def get_response_headers(_response_headers: dict | None = None) -> dict: +def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict: """ Sets the Appropriate OpenAI headers for the response and forward all headers as llm_provider-{header} @@ -31,7 +32,7 @@ def get_response_headers(_response_headers: dict | None = None) -> dict: return {**llm_provider_headers, **openai_headers} -def _get_llm_provider_headers(response_headers: dict) -> dict: +def _get_llm_provider_headers(response_headers: Mapping[str, str]) -> dict: """ Adds a llm_provider-{header} to all headers that are not already prefixed with llm_provider diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d9424d6a243..2b57883cc13 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -67,6 +67,96 @@ _BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$") _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$") _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$") _DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)") +_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:" +_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object]) +_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) + + +def is_claude_code_user_agent(user_agent: str) -> bool: + return user_agent.startswith("claude-cli/") + + +def _validated_claude_code_mapping(value: object) -> dict[object, object] | None: + try: + return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _validated_claude_code_list(value: object) -> list[object] | None: + try: + return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None: + stripped: Final = text.strip() + if "\n" in stripped or "\r" in stripped or not stripped.startswith(_CLAUDE_CODE_BILLING_HEADER_PREFIX): + return None + fields: Final = tuple( + field + for raw_field in stripped.removeprefix(_CLAUDE_CODE_BILLING_HEADER_PREFIX).split(";") + if (field := raw_field.strip()) + ) + if not fields or any("=" not in field for field in fields): + return None + parsed_fields: Final = tuple( + (parts[0].strip(), parts[1].strip()) for field in fields for parts in (field.split("=", 1),) + ) + if any(not key or not value for key, value in parsed_fields): + return None + return parsed_fields + + +def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None: + if isinstance(system, str): + return (system,) + blocks: Final = _validated_claude_code_list(system) + if blocks is None: + return None + block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks) + if any(block is None for block in block_mappings): + return None + text_values: Final = tuple( + block.get("text") for block in block_mappings if block is not None and block.get("type") == "text" + ) + if len(text_values) != len(blocks) or any(not isinstance(text, str) for text in text_values): + return None + meaningful_text: Final = tuple(text for text in text_values if isinstance(text, str) and text.strip()) + return meaningful_text or None + + +def _is_claude_code_subagent_billing_system(system: object) -> bool: + billing_texts: Final = _claude_code_billing_texts(system) + if billing_texts is None: + return False + billing_fields: Final = tuple( + fields for text in billing_texts if (fields := _claude_code_billing_fields(text)) is not None + ) + if len(billing_fields) != len(billing_texts): + return False + subagent_values: Final = tuple( + value for fields in billing_fields for key, value in fields if key == "cc_is_subagent" + ) + return subagent_values == ("true",) + + +def is_claude_code_one_shot_subagent_request( + messages: list[AllMessageValues], + system: object, + tools: object, + user_agent: str | None, +) -> bool: + only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None + return ( + user_agent is not None + and is_claude_code_user_agent(user_agent) + and not tools + and only_message is not None + and only_message.get("role") == "user" + and _is_claude_code_subagent_billing_system(system) + ) def _strip_bedrock_id_suffixes(model: str) -> str: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py index c1f10c245f8..d9cc65e730f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py @@ -82,7 +82,7 @@ async def anthropic_messages_with_mcp( LiteLLM_Proxy_MCP_Handler, ) - mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_references, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_references: return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 690040dd93b..6aa17372258 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -667,7 +667,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(response.read())) + raise BedrockError( + status_code=response.status_code, + message=str(response.read()), + headers=response.headers, + response=response, + ) # LOGGING logging_obj.post_call( @@ -690,6 +695,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( status_code=response.status_code, message=f"AgentCore: Failed to read/parse JSON response body: {e}", + headers=response.headers, ) parsed: Final = self._parse_json_response(response_json) @@ -880,7 +886,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(await response.aread())) + raise BedrockError( + status_code=response.status_code, + message=str(await response.aread()), + headers=response.headers, + response=response, + ) # LOGGING logging_obj.post_call( @@ -903,6 +914,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( status_code=response.status_code, message=f"AgentCore: Failed to read/parse JSON response body: {e}", + headers=response.headers, ) parsed: Final = self._parse_json_response(response_json) @@ -1031,6 +1043,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing response: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) def validate_environment( @@ -1046,7 +1059,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return headers def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) def should_fake_stream( self, diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index a75124325ae..984ba371898 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -22,7 +22,7 @@ from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token -from ..common_utils import BedrockError, _get_all_bedrock_regions +from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call @@ -66,7 +66,12 @@ def make_sync_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=str(response.read())) + raise BedrockError( + status_code=response.status_code, + message=str(response.read()), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -247,7 +252,12 @@ class BedrockConverseLLM(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -594,7 +604,12 @@ class BedrockConverseLLM(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index e097805f54a..fa24f8be893 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -2255,6 +2255,7 @@ class AmazonConverseConfig(BaseConfig): raise BedrockError( message=f"Error converting to valid response block={e}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues", status_code=422, + headers=response.headers, ) """ diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index e30ec731d8c..d489e47c3b5 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -470,6 +470,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing response: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) def validate_environment( @@ -485,7 +486,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): return headers def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) def should_fake_stream( self, diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index c39c88240c5..5f8a5544d65 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -42,6 +42,7 @@ from litellm.types.utils import GenericStreamingChunk as GChunk from ..common_utils import ( BedrockError, build_bedrock_stream_error, + error_response_text, get_bedrock_response_stream_shape, get_bedrock_tool_name, ) @@ -184,7 +185,12 @@ async def make_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=response.text) + raise BedrockError( + status_code=response.status_code, + message=error_response_text(response), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -228,9 +234,16 @@ async def make_call( ) return completion_stream, response.headers + except BedrockError: + raise except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") except Exception as e: @@ -270,7 +283,12 @@ def make_sync_call( ) if response.status_code != 200: - raise BedrockError(status_code=response.status_code, message=response.text) + raise BedrockError( + status_code=response.status_code, + message=error_response_text(response), + headers=response.headers, + response=response, + ) if fake_stream: model_response: Final[ModelResponse] = litellm.AmazonConverseConfig()._transform_response( @@ -314,9 +332,16 @@ def make_sync_call( ) return completion_stream, response.headers + except BedrockError: + raise except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=error_response_text(err.response), + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") except Exception as e: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index 04c6ec86a13..5d39b68d9d5 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -247,4 +247,4 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError: """Return the appropriate error class for Bedrock.""" - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index 1671585be2d..4bf1a1cba73 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -182,4 +182,4 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError: """Return the appropriate error class for Bedrock.""" - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index cd8066cda4d..d12c8aee48c 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -212,6 +212,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): raise BedrockError( message=f"Error parsing response: {raw_response.text}, error: {e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) verbose_logger.debug( @@ -241,6 +242,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): raise BedrockError( message=f"Error setting response content: {e}. Response: {completion_response}", status_code=raw_response.status_code, + headers=raw_response.headers, ) # Calculate usage from headers diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 37121d2ece7..a0e32c8aa22 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -295,7 +295,11 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): try: completion_response: Final = raw_response.json() except Exception: - raise BedrockError(message=raw_response.text, status_code=raw_response.status_code) + raise BedrockError( + message=raw_response.text, + status_code=raw_response.status_code, + headers=raw_response.headers, + ) verbose_logger.debug( "bedrock invoke response % s", json.dumps(completion_response, indent=4, default=str), @@ -363,6 +367,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error processing={raw_response.text}, Received error={e}", status_code=422, + headers=raw_response.headers, ) try: @@ -384,6 +389,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): raise BedrockError( message=f"Error parsing received text={outputText}.\nError-{e}", status_code=raw_response.status_code, + headers=raw_response.headers, ) ## CALCULATING USAGE - bedrock returns usage in the headers @@ -431,7 +437,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: - return BedrockError(status_code=status_code, message=error_message) + return BedrockError(status_code=status_code, message=error_message, headers=headers) @track_llm_api_timing() async def get_async_custom_stream_wrapper( diff --git a/litellm/llms/bedrock/claude_platform/common_utils.py b/litellm/llms/bedrock/claude_platform/common_utils.py index 311f3a56b84..fb7f2185ec5 100644 --- a/litellm/llms/bedrock/claude_platform/common_utils.py +++ b/litellm/llms/bedrock/claude_platform/common_utils.py @@ -1,7 +1,10 @@ from typing import Final +import httpx + import litellm from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.secret_managers.main import get_secret_str CLAUDE_PLATFORM_SERVICE_NAME: Final = "aws-external-anthropic" @@ -15,6 +18,14 @@ def strip_claude_platform_route(model: str) -> str: class BedrockClaudePlatformMixin(BaseAWSLLM): + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + @staticmethod def _get_workspace_id(optional_params: dict, litellm_params: dict) -> str | None: workspace_id = ( diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index fe675a30a00..be4f0f32689 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -33,8 +33,53 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues +_ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs" + + +def error_response_text(response: httpx.Response) -> str: + try: + return response.text + except httpx.ResponseNotRead: + return response.reason_phrase + + +def _synthesize_error_response( + *, status_code: int, headers: dict[str, object] | httpx.Headers, request: httpx.Request | None +) -> tuple[httpx.Request, httpx.Response]: + error_request: Final = request or httpx.Request(method="POST", url=_ERROR_REQUEST_URL) + safe_headers: Final = ( + headers + if isinstance(headers, httpx.Headers) + else tuple((key, value) for key, value in headers.items() if isinstance(value, (str, bytes))) + ) + return error_request, httpx.Response(status_code=status_code, headers=safe_headers, request=error_request) + + class BedrockError(BaseLLMException): - pass + def __init__( + self, + status_code: int, + message: str, + headers: dict[str, object] | httpx.Headers | None = None, + request: httpx.Request | None = None, + response: httpx.Response | None = None, + body: dict[str, object] | None = None, + status_code_is_synthesized: bool = False, + ) -> None: + error_request, error_response = ( + _synthesize_error_response(status_code=status_code, headers=headers, request=request) + if response is None and headers + else (request, response) + ) + super().__init__( + status_code=status_code, + message=message, + headers=headers, + request=error_request, + response=error_response, + body=body, + status_code_is_synthesized=status_code_is_synthesized, + ) _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = ( diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index 2383350b3a3..1fb53f6ff0a 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -102,6 +102,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise BedrockError( status_code=response.status_code, message=error_text, + headers=response.headers, + response=response, ) bedrock_response: Final = response.json() @@ -124,6 +126,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig): raise BedrockError( status_code=e.response.status_code, message=e.response.text, + headers=e.response.headers, + response=e.response, ) except Exception as e: verbose_logger.error("Error in CountTokens handler: %s", e) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 5fb86d476f4..d3725434498 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -132,7 +132,12 @@ class BedrockEmbedding(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -161,7 +166,12 @@ class BedrockEmbedding(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py index 18d47301ee5..acb0cc8dcb7 100644 --- a/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py +++ b/litellm/llms/bedrock/image_edit/amazon_nova_canvas_image_edit_transformation.py @@ -20,6 +20,7 @@ import httpx from litellm._logging import verbose_logger from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import FileTypes, ImageObject, ImageResponse @@ -228,6 +229,14 @@ class BedrockAmazonNovaCanvasImageEditConfig(BaseImageEditConfig): """ return _supports_nova_canvas_image_edit_from_model_cost(model or "") + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_supported_openai_params(self, model: str) -> list: return [ "n", diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 5c517f2049c..be6489f20ae 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -114,7 +114,12 @@ class BedrockImageEdit(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -156,7 +161,12 @@ class BedrockImageEdit(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/image_edit/stability_transformation.py b/litellm/llms/bedrock/image_edit/stability_transformation.py index 24e7ba73075..bc9a64f587a 100644 --- a/litellm/llms/bedrock/image_edit/stability_transformation.py +++ b/litellm/llms/bedrock/image_edit/stability_transformation.py @@ -27,6 +27,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.llms.stability import ( OPENAI_SIZE_TO_STABILITY_ASPECT_RATIO, @@ -84,6 +85,14 @@ class BedrockStabilityImageEditConfig(BaseImageEditConfig): return True return False + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_supported_openai_params(self, model: str) -> list: """ Return list of OpenAI params supported by Bedrock Stability. diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index c78e3c147cb..87762b648e0 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -119,7 +119,12 @@ class BedrockImageGeneration(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") ### FORMAT RESPONSE TO OPENAI FORMAT ### @@ -162,7 +167,12 @@ class BedrockImageGeneration(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 6ff9f0155f9..a715d150b4c 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -29,6 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( + BedrockError, apply_bedrock_invoke_structured_output, ensure_bedrock_anthropic_messages_tool_names, get_anthropic_beta_from_headers, @@ -79,6 +80,14 @@ class AmazonAnthropicClaudeMessagesConfig( BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS = frozenset(BedrockInvokeAnthropicMessagesRequest.__annotations__.keys()) + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def __init__(self, **kwargs): BaseAnthropicMessagesConfig.__init__(self, **kwargs) AmazonInvokeConfig.__init__(self, **kwargs) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index d0a3c37ffb3..fb8bc4f191f 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -2,13 +2,14 @@ import json from collections.abc import Mapping from typing import TYPE_CHECKING, Final, Optional, cast +import httpx from httpx import Response from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockEventStreamDecoderBase, BedrockModelInfo +from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo if TYPE_CHECKING: from httpx import URL @@ -18,6 +19,14 @@ if TYPE_CHECKING: class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamDecoderBase, BasePassthroughConfig): + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def is_streaming_request(self, endpoint: str, request_data: dict) -> bool: return "stream" in endpoint diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 1f4c81d6491..3b972961940 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -9,12 +9,14 @@ import json import uuid as uuid_lib from typing import Final, cast +import httpx from pydantic import BaseModel from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm from litellm.types.llms.openai import ( OpenAIRealtimeContentPartDone, @@ -121,6 +123,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self._cumulative_usage = BedrockUsageEvent() self._reported_usage = BedrockUsageEvent() + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def validate_environment(self, headers: dict, model: str, api_key: str | None = None) -> dict: """Validate environment - no special validation needed for Bedrock.""" return headers diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 4860c99268e..8847381cbc9 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -46,7 +46,12 @@ class BedrockRerankHandler(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") @@ -117,7 +122,12 @@ class BedrockRerankHandler(BaseAWSLLM): response.raise_for_status() except httpx.HTTPStatusError as err: error_code: Final = err.response.status_code - raise BedrockError(status_code=error_code, message=err.response.text) + raise BedrockError( + status_code=error_code, + message=err.response.text, + headers=err.response.headers, + response=err.response, + ) except httpx.TimeoutException: raise BedrockError(status_code=408, message="Timeout error occurred.") diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 920e566c9dd..e7d706c3731 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -39,7 +39,6 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -380,6 +379,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, message=f"AgentCore gateway MCP error: {error}", + headers=raw_response.headers, ) # A failed tools/call is reported in-band, as HTTP 200 with result.isError @@ -389,6 +389,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=raw_response.status_code if raw_response.status_code >= 400 else 502, message=f"AgentCore web search tool error: {self._tool_error_message(response_json)}", + headers=raw_response.headers, ) text_items: Final = tuple( @@ -440,6 +441,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): raise BedrockError( status_code=502, message=f"AgentCore gateway returned SSE without a JSON data frame: {text[:200]}", + headers=raw_response.headers, ) def get_error_class( @@ -448,7 +450,7 @@ class AgentCoreSearchConfig(BaseSearchConfig, BaseAWSLLM): status_code: int, headers: dict, # mutable-ok: BaseSearchConfig.get_error_class takes the response headers as a dict ) -> Exception: - return BaseLLMException( + return BedrockError( status_code=status_code, message=error_message, headers=headers, diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 6940077391f..27c90c9d71e 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBContent, BedrockKBResponse, @@ -38,6 +39,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): BaseVectorStoreConfig.__init__(self) BaseAWSLLM.__init__(self) + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, # mutable-ok: base passes response headers as a dict + ) -> BedrockError: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials: return {} diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 64d7ef2bed6..d91157c3d10 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -13,6 +13,8 @@ Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the from collections.abc import AsyncIterator, Iterator from typing import Any, Final +import httpx + import litellm from litellm._logging import verbose_logger from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -24,6 +26,8 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams +from ...base_llm.chat.transformation import BaseLLMException +from ...bedrock.common_utils import BedrockError from ...openai_like.chat.transformation import OpenAILikeChatConfig from ..common_utils import mantle_base_segment @@ -45,6 +49,11 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): def get_config(cls): return super().get_config() + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers + ) -> BaseLLMException: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def _get_openai_compatible_provider_info( self, api_base: str | None, diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 5179c966584..bbbda4d14b6 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -19,11 +19,14 @@ import json from collections.abc import Mapping from typing import Any, Final +import httpx from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock_mantle.common_utils import ( MANTLE_HOST_RE, BedrockMantleAuthMixin, @@ -98,6 +101,11 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK_MANTLE + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers + ) -> BaseLLMException: + return BedrockError(status_code=status_code, message=error_message, headers=headers) + def get_complete_url( self, api_base: str | None, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2f561809940..8b6d3a0ceb4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + return self._transform_ocr_response( provider_config=provider_config, model=model, @@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + # Use async response transform for async operations return await provider_config.async_transform_ocr_response( model=model, diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index ac934ad0cb5..21a630a76d7 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final from httpx import Headers @@ -13,7 +15,7 @@ class FireworksAIException(BaseLLMException): pass -def get_fireworks_session_id(litellm_params: dict) -> str | None: +def get_fireworks_session_id(litellm_params: Mapping[str, object]) -> str | None: """ Session id to send as `x-session-affinity`, or None when the caller gave none. @@ -23,19 +25,39 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None: """ params: Final = litellm_params metadata: Final = params.get("metadata") - if isinstance(metadata, dict) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): + if isinstance(metadata, Mapping) and metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None for key in ("litellm_session_id", "session_id"): value = params.get(key) if value: return str(value) - if isinstance(metadata, dict): + if isinstance(metadata, Mapping): value = metadata.get("session_id") if value: return str(value) return None +def with_fireworks_session_affinity( + headers: Mapping[str, str], litellm_params: Mapping[str, object] +) -> Mapping[str, str]: + if any(key.lower() == "x-session-affinity" for key in headers): + return headers + session_id: Final = get_fireworks_session_id(litellm_params) + if not session_id: + return headers + return MappingProxyType({**headers, "x-session-affinity": session_id}) + + +def resolve_fireworks_api_key(api_key: str | None) -> str | None: + return api_key or ( + get_secret_str("FIREWORKS_API_KEY") + or get_secret_str("FIREWORKS_AI_API_KEY") + or get_secret_str("FIREWORKSAI_API_KEY") + or get_secret_str("FIREWORKS_AI_TOKEN") + ) + + AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-" @@ -63,13 +85,7 @@ class FireworksAIMixin: ) def _get_api_key(self, api_key: str | None) -> str | None: - dynamic_api_key: Final = api_key or ( - get_secret_str("FIREWORKS_API_KEY") - or get_secret_str("FIREWORKS_AI_API_KEY") - or get_secret_str("FIREWORKSAI_API_KEY") - or get_secret_str("FIREWORKS_AI_TOKEN") - ) - return dynamic_api_key + return resolve_fireworks_api_key(api_key) def validate_environment( self, @@ -92,9 +108,5 @@ class FireworksAIMixin: return self._add_session_affinity_header({**auth_headers, **content_type_header}, litellm_params) def _add_session_affinity_header(self, headers: dict, litellm_params: dict) -> dict: - if any(key.lower() == "x-session-affinity" for key in headers): - return headers - session_id: Final = get_fireworks_session_id(litellm_params) - if not session_id: - return headers - return {**headers, "x-session-affinity": session_id} + pinned: Final = with_fireworks_session_affinity(headers, litellm_params) + return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py new file mode 100644 index 00000000000..fb0587553d4 --- /dev/null +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -0,0 +1,102 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final +from urllib.parse import unquote + +import httpx +from openai.types.responses import EasyInputMessageParam, ResponseInputItemParam + +from litellm.llms.fireworks_ai.common_utils import ( + resolve_fireworks_api_key, + resolve_fireworks_resource_name, + with_fireworks_session_affinity, +) +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ResponseInputParam +from litellm.types.responses.main import DeleteResponseResult +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +FIREWORKS_AI_DEFAULT_API_BASE: Final = "https://api.fireworks.ai/inference/v1" + + +def _session_params(litellm_params: GenericLiteLLMParams) -> Mapping[str, object]: + extras: Final[Mapping[str, object]] = litellm_params.model_extra or MappingProxyType({}) + return MappingProxyType( + {"litellm_session_id": extras.get("litellm_session_id"), "metadata": extras.get("litellm_metadata")} + ) + + +def _developer_item_as_system(item: ResponseInputItemParam) -> ResponseInputItemParam: + if "role" not in item or item["role"] != "developer": + return item + return EasyInputMessageParam(role="system", content=item["content"], type="message") + + +def _developer_items_as_system(input: str | ResponseInputParam) -> str | ResponseInputParam: + if isinstance(input, str): + return input + return [_developer_item_as_system(item) for item in input] + + +class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): + @property + def custom_llm_provider(self) -> LlmProviders: + return LlmProviders.FIREWORKS_AI + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + litellm_params: GenericLiteLLMParams | None, + ) -> dict: # mutable-ok: overrides the base class signature + params: Final = litellm_params or GenericLiteLLMParams() + api_key: Final = resolve_fireworks_api_key(params.api_key) + if api_key is None: + raise ValueError("FIREWORKS_API_KEY is not set") + authorized: Final = MappingProxyType( + {"Content-Type": "application/json", **headers, "Authorization": f"Bearer {api_key}"} + ) + pinned: Final = with_fireworks_session_affinity(authorized, _session_params(params)) + return dict(pinned) # mutable-ok: the HTTP handler updates the returned headers in place + + def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/") + return f"{base}/responses" + + def _validate_input_param(self, input: str | ResponseInputParam) -> str | ResponseInputParam: + return _developer_items_as_system(super()._validate_input_param(input)) + + def transform_responses_api_request( + self, + model: str, + input: str | ResponseInputParam, + response_api_optional_request_params: dict, # mutable-ok: overrides the base class signature + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: overrides the base class signature + ) -> dict: # mutable-ok: overrides the base class signature + return super().transform_responses_api_request( + model=resolve_fireworks_resource_name(model), + input=input, + response_api_optional_request_params=response_api_optional_request_params, + litellm_params=litellm_params, + headers=headers, + ) + + def transform_delete_response_api_response( + self, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> DeleteResponseResult: + deleted_id: Final = unquote(raw_response.request.url.path.rsplit("/", 1)[-1]) + return DeleteResponseResult(id=deleted_id, object="response", deleted=True) + + def supports_native_websocket(self) -> bool: + return False + + def supports_native_file_search(self) -> bool: + return False 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/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/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 377cd9f3437..a15ea4d845b 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,6 +1,7 @@ import json -from collections.abc import Coroutine +from collections.abc import Coroutine, Sequence from typing import TYPE_CHECKING, Final, Protocol +from urllib.parse import urlparse import httpx from typing_extensions import ReadOnly, TypedDict @@ -12,11 +13,13 @@ from litellm.litellm_core_utils.url_utils import ( safe_get, ) from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler, _get_httpx_client, get_async_httpx_client, ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( VERTEX_CREDENTIALS_TYPES, @@ -55,6 +58,20 @@ class _FetchedResponseView(TypedDict): response: ReadOnly[httpx.Response] +class _VertexEndpointDeployedModel(TypedDict, total=False): + model: ReadOnly[str] + + +class _VertexEndpointResponse(TypedDict, total=False): + deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]] + + +class _VertexEndpointPayloadView(TypedDict): + """Holds one decoded GET endpoints/ response so the payload reads back typed.""" + + payload: ReadOnly[_VertexEndpointResponse] + + def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse: return response.json() @@ -78,7 +95,17 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location: str | None, timeout: float | httpx.Timeout, max_retries: int | None, + custom_endpoint: bool | None = None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: + if custom_endpoint: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. " + "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; " + "use a publisher model or fine-tuned Gemini endpoint deployment instead." + ), + ) sync_handler: Final = _get_httpx_client() access_token, project_id = self._ensure_access_token( @@ -87,6 +114,26 @@ class VertexAIBatchPrediction(VertexLLM): custom_llm_provider="vertex_ai", ) + headers: Final = { + "Content-Type": "application/json; charset=utf-8", + "Authorization": f"Bearer {access_token}", + } + + transformed_batch_request: Final[VertexAIBatchPredictionJob] = ( + VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( + request=create_batch_data, + vertex_project=vertex_project or project_id, + vertex_location=vertex_location or "us-central1", + ) + ) + vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model( + vertex_batch_request=transformed_batch_request, + headers=headers, + sync_handler=sync_handler, + api_base=api_base, + vertex_location=vertex_location or "us-central1", + ) + default_api_base: Final = self.create_vertex_batch_url( vertex_location=vertex_location or "us-central1", vertex_project=vertex_project or project_id, @@ -111,17 +158,6 @@ class VertexAIBatchPrediction(VertexLLM): vertex_api_version="v1", ) - headers: Final = { - "Content-Type": "application/json; charset=utf-8", - "Authorization": f"Bearer {access_token}", - } - - vertex_batch_request: Final[VertexAIBatchPredictionJob] = ( - VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( - request=create_batch_data - ) - ) - if _is_async is True: return self._async_create_batch( vertex_batch_request=vertex_batch_request, @@ -142,6 +178,77 @@ class VertexAIBatchPrediction(VertexLLM): ) return vertex_batch_response + @staticmethod + def _build_endpoint_resolution_url(api_base: str | None, model: str, vertex_location: str) -> str: + """ + Builds the GET url for resolving an endpoint resource (`projects/../endpoints/`). + + A custom `api_base` replaces the Google host: its `/v1`/`/v1beta1` path swallows the + version segment (matching `_check_custom_proxy`'s grafting), any other path is kept as a + mount prefix in front of the full default path. The `:operation` suffix convention from + `_check_custom_proxy` does not apply to a plain resource GET. + """ + default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}" + if not api_base: + return default_endpoint_url + api_base_path: Final = urlparse(api_base).path.rstrip("/") + if api_base_path in ("/v1", "/v1beta1"): + return _graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url) + return api_base.rstrip("/") + urlparse(default_endpoint_url).path + + def _resolve_fine_tuned_endpoint_model( + self, + vertex_batch_request: VertexAIBatchPredictionJob, + headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers + sync_handler: HTTPHandler, + api_base: str | None, + vertex_location: str, + ) -> VertexAIBatchPredictionJob: + """ + A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only + accepts Model resources, so swap the endpoint resource for its deployed tuned model + (`projects/../locations/../models/`) read from GET endpoints/. + """ + model: Final = vertex_batch_request.get("model", "") + if "/endpoints/" not in model: + return vertex_batch_request + + endpoint_url: Final = self._build_endpoint_resolution_url( + api_base=api_base, + model=model, + vertex_location=vertex_location, + ) + # ``api_base`` can come from caller-supplied request kwargs, so wrap the + # fetch in ``safe_get``: it rejects DNS-rebind / private / cloud-metadata + # targets before the bearer token leaves the process (mirrors retrieve_batch). + fetched: Final[_FetchedResponseView] = { + "response": safe_get( + sync_handler, + endpoint_url, + headers=headers, + ) + } + response: Final = fetched["response"] + if response.status_code != 200: + raise VertexAIError( + status_code=response.status_code, + message=f"Failed to resolve fine-tuned Vertex endpoint '{model}': {response.text}", + ) + + payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()} + deployed_models: Final = payload_view["payload"].get("deployedModels") or () + deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else "" + if not deployed_model: + raise VertexAIError( + status_code=400, + message=( + f"Vertex endpoint '{model}' has no deployed model, so there is no tuned model " + "resource to run batch predictions against" + ), + ) + resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model} + return resolved_request + async def _async_create_batch( self, vertex_batch_request: VertexAIBatchPredictionJob, diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index f284b47292b..e63c80dd3cf 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -22,6 +22,8 @@ class VertexAIBatchTransformation: def transform_openai_batch_request_to_vertex_ai_batch_request( cls, request: CreateBatchRequest, + vertex_project: str | None = None, + vertex_location: str | None = None, ) -> VertexAIBatchPredictionJob: """ Transforms OpenAI Batch requests to Vertex AI Batch requests @@ -31,7 +33,11 @@ class VertexAIBatchTransformation: if input_file_id is None: raise ValueError("input_file_id is required, but not provided") input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl") - model: Final[str] = cls._get_model_from_gcs_file(input_file_id) + model: Final[str] = cls._get_batch_job_model( + input_file_id=input_file_id, + vertex_project=vertex_project, + vertex_location=vertex_location, + ) output_config: Final[OutputConfig] = OutputConfig( predictionsFormat="jsonl", gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)), @@ -188,6 +194,33 @@ class VertexAIBatchTransformation: path_parts: Final = input_file_id.rsplit("/", 1) return path_parts[0] + @classmethod + def _get_batch_job_model( + cls, + input_file_id: str, + vertex_project: str | None, + vertex_location: str | None, + ) -> str: + """ + Returns the `model` for the batchPredictionJobs request: the publisher model path as-is, or + the full `projects/../locations/../endpoints/` resource name for a fine-tuned endpoint. + + The v1 batch API only accepts Model resources, so the handler resolves an endpoint resource + to its deployed tuned model (`projects/../locations/../models/`) before sending the job. + """ + parsed_model: Final = cls._get_model_from_gcs_file(input_file_id) + if not parsed_model.startswith("endpoints/"): + return parsed_model + if not vertex_project: + raise VertexAIError( + status_code=400, + message=( + f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require " + "`vertex_project` to build the endpoint resource name" + ), + ) + return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}" + @classmethod def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str: """ @@ -202,6 +235,9 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + Fine-tuned Gemini endpoints are stored as `endpoints/` in the uri and returned + in that form. + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ model: Final = cls._parse_model_from_gcs_file(gcs_file_uri) @@ -210,11 +246,13 @@ class VertexAIBatchTransformation: status_code=400, message=( "Vertex AI batch creation requires the model to be part of `input_file_id`, but " - f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + f"'{gcs_file_uri}' contains no 'publishers//models/' or " + "'endpoints/' path segment. " "Either upload the input file through LiteLLM (POST /v1/files with " "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " "pass a uri of the form " - "gs:////publishers//models//" + "gs:////publishers//models// " + "(or gs:////endpoints// for fine-tuned models)" ), ) return model @@ -222,18 +260,26 @@ class VertexAIBatchTransformation: @classmethod def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: """ - Returns the `publishers//models/` path from a gcs uri, or None if the uri - does not contain one. + Returns the `publishers//models/` or `endpoints/` path from a + gcs uri, or None if the uri does not contain one. + + A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is + used, so a user-configured bucket prefix that happens to contain `endpoints/` cannot + override the model path LiteLLM appended after it. """ - _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") - if not separator: - return None + unquoted_uri: Final = unquote(gcs_file_uri) + _, separator, model_path = unquoted_uri.partition("publishers/") + if separator: + parts: Final = model_path.split("/") + if len(parts) >= 3 and parts[1] == "models" and parts[2]: + return f"publishers/{'/'.join(parts[:3])}" - parts: Final = model_path.split("/") - if len(parts) < 3 or parts[1] != "models" or not parts[2]: - return None + _, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/") + endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else "" + if endpoint_id.isdigit(): + return f"endpoints/{endpoint_id}" - return f"publishers/{'/'.join(parts[:3])}" + return None @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index fe2e0ab6c06..14aebcaabaf 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -370,6 +370,19 @@ def get_vertex_base_model_name(model: str) -> str: return model +def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None: + """ + Fine-tuned Gemini deployments are addressed by a numeric endpoint id, + configured as `vertex_ai/` or `vertex_ai/gemini/`. + + Returns the endpoint id, or None when `model` is a regular publisher model. + Mirrors the online chat path in `_get_vertex_url`, which sends numeric + models to `endpoints/{id}` instead of `publishers/google/models/{model}`. + """ + candidate: Final = model.split("/")[-1] if "gemini/" in model else model + return candidate if candidate.isdigit() else None + + def validate_vertex_location(vertex_location: str | None) -> str: """ Validate a Vertex AI location before interpolating it into a request host or diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index b6ad9fbcc04..263956efc9f 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -39,6 +39,7 @@ from litellm.llms.base_llm.files.transformation import ( ) from litellm.llms.vertex_ai.common_utils import ( _convert_vertex_datetime_to_openai_datetime, + get_vertex_ai_fine_tuned_endpoint_id, ) from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -707,20 +708,39 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, openai_jsonl_content: list[dict[str, Any]], + deployment_model: str | None = None, ) -> str: """ Gets a unique GCS object name for the VertexAI batch prediction job named as: litellm-vertex-{model}-{uuid} + + The stored model path decides which Vertex model the batch job later executes against, so + `deployment_model` (the deployment's own configured model) wins over the user-supplied + JSONL `body.model`; the JSONL value is only a fallback for direct SDK calls that carry no + deployment config. + + Fine-tuned Gemini deployments (numeric endpoint ids) are stored under + `endpoints/` so the batch transformation can round-trip them into a + `projects/../locations/../endpoints/` batch job model instead of a + nonexistent publisher model. """ - _model = openai_jsonl_content[0].get("body", {}).get("model", "") - if "publishers/google/models" not in _model: - _model = f"publishers/google/models/{_model}" - safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model") + raw_model: Final = ( + deployment_model.removeprefix("vertex_ai/") + if deployment_model + else openai_jsonl_content[0].get("body", {}).get("model", "") + ) + endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) + model_path: Final = ( + f"endpoints/{endpoint_id}" + if endpoint_id is not None + else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") + ) + safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name - def get_object_name(self, file_data: FileTypes, purpose: str) -> str: + def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str: """ Get the object name for the request. @@ -728,10 +748,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): upload is never materialized just to derive the GCS object name. """ if purpose == "batch": - ## 1. If jsonl, derive the object name from the first entry's model + ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) if first_entry is not None: - return self._get_gcs_object_name_from_batch_jsonl([first_entry]) + return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model) ## 2. If not jsonl, store under a server-generated managed object name filename, _ = extract_file_metadata(file_data) @@ -761,6 +781,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Get the complete url for the request """ + if data.get("purpose") == "batch" and litellm_params.get("custom_endpoint"): + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. " + "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; " + "remove this deployment from the batch request (e.g. `target_model_names`) or " + "use a publisher model / fine-tuned Gemini endpoint instead." + ), + ) bucket_name = self._get_configured_bucket_name(litellm_params) bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) file_data: Final = data.get("file") @@ -769,7 +799,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - object_name = self.get_object_name(file_data, purpose) + configured_model: Final = litellm_params.get("model") + object_name = self.get_object_name( + file_data, + purpose, + deployment_model=configured_model if isinstance(configured_model, str) else None, + ) if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name: Final = encode_gcs_object_name_for_url(object_name) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 32563c32aa1..b1ffc1583e4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -30791,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", @@ -30808,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, @@ -37225,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, @@ -37265,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, @@ -37476,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, @@ -37516,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, diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index 6bf21a19896..7ccff9434a7 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -6,10 +6,12 @@ Canonical definition for ``litellm_mcpservertable``. Re-exported from """ import enum +from collections.abc import Mapping from datetime import datetime +from types import MappingProxyType from typing import Literal -from pydantic import Field +from pydantic import Field, ValidationInfo, field_validator from litellm.types.llms.base import LiteLLMPydanticObjectBase from litellm.types.mcp import MCPAuthType, MCPCredentials, MCPTransportType @@ -115,3 +117,12 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): submitted_at: datetime | None = None reviewed_at: datetime | None = None review_notes: str | None = None + + @field_validator("static_headers", "env", mode="before") + @classmethod + def decode_stored_secret_map(cls, value: object, info: ValidationInfo) -> Mapping[str, str] | None: + from litellm.proxy.common_utils.encrypt_decrypt_utils import decode_secret_map + + if value is None and info.field_name == "env": + return MappingProxyType({}) + return decode_secret_map(value, key=info.field_name or "secret map") diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index ad9622c9cd0..4bac65125b4 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -3108,15 +3108,17 @@ class MCPRequestHandler: @staticmethod async def _get_allowed_mcp_servers_for_agent( user_api_key_auth: UserAPIKeyAuth | None = None, - agent_object_permission=None, + agent_object_permission: LiteLLM_ObjectPermissionTable | None = None, ) -> list[str]: """ Get allowed MCP servers for an agent (from the agent's object_permission). - Returns the MCP servers from the agent's object_permission. - If agent has no object_permission, returns [] (no extra restriction). An entitlement the - agent LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here so the - resolver denies. + Returns the agent's direct servers, the servers in its access groups, and the servers reached + through its toolsets, exactly as the key, team, and org levels count theirs. If agent has no + object_permission, returns [] (no extra restriction). An entitlement the agent LINKS but that + cannot be read, or a declared toolset that resolves to no grants, raises + ``UnloadableEntitlementError`` out of here so the resolver denies instead of reading the + agent as unrestricted. Args: user_api_key_auth: User auth with agent_id @@ -3126,31 +3128,30 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return [] - obj_perm = agent_object_permission - if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + obj_perm: Final = ( + agent_object_permission + if agent_object_permission is not None + else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + ) if obj_perm is None: return [] try: - direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or [] - if isinstance(direct_mcp_servers, str): - direct_mcp_servers = [] - mcp_access_groups = getattr(obj_perm, "mcp_access_groups", None) or [] - if isinstance(mcp_access_groups, str): - mcp_access_groups = [] - - # Permission entries may be server_ids OR names/aliases — expand to ids. from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list(list(direct_mcp_servers)) - - access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups(mcp_access_groups) - all_servers: Final = expanded_direct_servers + access_group_servers - return list(set(all_servers)) + expanded_direct_servers: Final = global_mcp_server_manager.expand_permission_list( + obj_perm.mcp_servers or [] + ) + access_group_servers: Final = await MCPRequestHandler._get_mcp_servers_from_access_groups( + obj_perm.mcp_access_groups or [] + ) + toolset_grants: Final = await MCPRequestHandler._toolset_tool_permissions(obj_perm) + return list({*expanded_direct_servers, *access_group_servers, *toolset_grants}) except Exception as e: + if isinstance(e, UnloadableEntitlementError): + raise verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e) return [] @@ -3158,13 +3159,15 @@ class MCPRequestHandler: async def _get_agent_tool_permissions_for_server( server_id: str, user_api_key_auth: UserAPIKeyAuth | None = None, - agent_object_permission=None, + agent_object_permission: LiteLLM_ObjectPermissionTable | None = None, ) -> list[str] | None: """ - Get allowed tool names for a server from the agent's object_permission. - Returns None if agent has no tool restrictions for this server. An entitlement the agent - LINKS but that cannot be read raises ``UnloadableEntitlementError`` out of here, which the - tool resolver turns into deny-all for the server rather than an unrestricted tool list. + Get allowed tool names for a server from the agent's object_permission: the union of its + direct tool permissions and the tools its toolsets grant on that server, mirroring the key and + team levels. Returns None if agent has no tool restrictions for this server. An entitlement the + agent LINKS but that cannot be read, or a declared toolset that resolves to no grants, raises + ``UnloadableEntitlementError`` out of here, which the tool resolver turns into deny-all for the + server rather than an unrestricted tool list. Args: server_id: Server ID to check permissions for @@ -3175,24 +3178,30 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.agent_id: return None - obj_perm = agent_object_permission - if obj_perm is None: - obj_perm = await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + obj_perm: Final = ( + agent_object_permission + if agent_object_permission is not None + else await MCPRequestHandler._get_agent_object_permission(user_api_key_auth) + ) if obj_perm is None: return None try: - mcp_tool_permissions: Final = getattr(obj_perm, "mcp_tool_permissions", None) - if not mcp_tool_permissions or not isinstance(mcp_tool_permissions, dict): - return None - # Dict keys may be server_ids OR names/aliases; normalize before lookup. from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - tools: Final = global_mcp_server_manager.expand_tool_permissions(mcp_tool_permissions).get(server_id) - return list(tools) if tools else None + direct_tools: Final = ( + global_mcp_server_manager.expand_tool_permissions(obj_perm.mcp_tool_permissions).get(server_id) + if obj_perm.mcp_tool_permissions + else None + ) + toolset_tools: Final = await MCPRequestHandler._toolset_tools_for_server(obj_perm, server_id) + agent_tools: Final = MCPRequestHandler._union_tool_grants(direct_tools, toolset_tools) + return list(agent_tools) if agent_tools else None except Exception as e: + if isinstance(e, UnloadableEntitlementError): + raise verbose_logger.warning("Failed to get agent tool permissions for server: %s", e) return None diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 41d0b78b555..7379126983a 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -4,7 +4,7 @@ import hashlib import json from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -23,8 +23,11 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + SecretMapDecodeError, _get_salt_key, + decode_secret_map, decrypt_value_helper, + encrypt_secret_map, encrypt_value_helper, ) from litellm.proxy.utils import PrismaClient @@ -122,6 +125,9 @@ class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False): server_id: str +OAuthGrantState = Literal["valid", "refreshable", "absent"] + + class _OAuthTokenRefreshResponse(TypedDict, total=False): access_token: str refresh_token: str @@ -360,7 +366,7 @@ def _prepare_mcp_server_data( # exclude_unset filter is respected. Reading back from ``data`` would # reintroduce defaults (e.g. ``env={}``) for fields the caller never set. if data_dict.get("static_headers") is not None: - data_dict["static_headers"] = safe_dumps(data_dict["static_headers"]) + data_dict["static_headers"] = encrypt_secret_map(data_dict["static_headers"]) # env_vars is read from ``data_dict`` (not ``data``) like every other JSON # column so the exclude_unset filter is respected: a partial update that @@ -376,7 +382,7 @@ def _prepare_mcp_server_data( data_dict["mcp_info"] = safe_dumps(data_dict["mcp_info"]) if data_dict.get("env") is not None: - data_dict["env"] = safe_dumps(data_dict["env"]) + data_dict["env"] = encrypt_secret_map(data_dict["env"]) if "tool_name_to_display_name" in data_dict: data_dict["tool_name_to_display_name"] = safe_dumps(data_dict["tool_name_to_display_name"] or {}) @@ -589,6 +595,19 @@ def decrypt_credentials( return credentials +def _readable_mcp_servers( + rows: Iterable["prisma_db_models.LiteLLM_MCPServerTable"], +) -> Iterable[LiteLLM_MCPServerTable]: + for row in rows: + try: + table = LiteLLM_MCPServerTable.model_validate(row.model_dump()) + except SecretMapDecodeError: + verbose_proxy_logger.warning("Skipping MCP server %s: cannot decrypt secret map", row.server_id) + continue + decrypt_global_env_var_values(table.env_vars) + yield table + + async def get_all_mcp_servers( prisma_client: PrismaClient, approval_status: str | None = None, @@ -609,10 +628,7 @@ async def get_all_mcp_servers( ) mcp_servers: Final = await _db_find_mcp_server_rows(prisma_client, where) - tables: Final = [LiteLLM_MCPServerTable.model_validate(mcp_server.model_dump()) for mcp_server in mcp_servers] - for table in tables: - decrypt_global_env_var_values(table.env_vars) - return tables + return list(_readable_mcp_servers(mcp_servers)) async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> LiteLLM_MCPServerTable | None: @@ -638,13 +654,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str] "server_id": {"in": server_ids}, } ) - final_mcp_servers: Final[list[LiteLLM_MCPServerTable]] = [] - for _mcp_server in _mcp_servers: - table = LiteLLM_MCPServerTable.model_validate(_mcp_server.model_dump()) - decrypt_global_env_var_values(table.env_vars) - final_mcp_servers.append(table) - - return final_mcp_servers + return list(_readable_mcp_servers(_mcp_servers)) async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> list[str]: @@ -852,12 +862,10 @@ async def create_mcp_server( data_dict["created_by"] = touched_by data_dict["updated_by"] = touched_by - new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create( - data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable - ) + new_mcp_server: Final = await MCPServerRepository(prisma_client).table.create(data=data_dict) _decrypt_env_vars_on_returned_row(new_mcp_server) - return new_mcp_server + return LiteLLM_MCPServerTable.model_validate(new_mcp_server.model_dump()) async def create_draft_mcp_server( @@ -1066,13 +1074,13 @@ async def update_mcp_server( data_dict["credentials"] = Json(None) - updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update( + updated_mcp_server: Final = await MCPServerRepository(prisma_client).table.update( where={"server_id": data.server_id}, - data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable + data=data_dict, ) _decrypt_env_vars_on_returned_row(updated_mcp_server) - return updated_mcp_server + return LiteLLM_MCPServerTable.model_validate(updated_mcp_server.model_dump()) if updated_mcp_server else None async def get_mcp_server_oauth_client_credentials(prisma_client: PrismaClient, server_id: str) -> object | None: @@ -1144,6 +1152,13 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient, if rotated_env_vars is not None: update_data["env_vars"] = safe_dumps(rotated_env_vars) + for field in ("static_headers", "env"): + try: + if secret_map := decode_secret_map(getattr(mcp_server, field, None), key=field): + update_data[field] = encrypt_secret_map(secret_map, new_encryption_key=new_master_key) + except SecretMapDecodeError: + verbose_proxy_logger.warning("Cannot rotate MCP %s for server %s", field, mcp_server.server_id) + if not update_data: continue @@ -1453,6 +1468,15 @@ def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: in return False +def oauth_grant_state(cred: OAuthCredentialPayload | None) -> OAuthGrantState: + """Classify local grant readiness without attempting a refresh or checking upstream revocation.""" + if not cred or not cred.get("access_token"): + return "absent" + if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): + return "valid" + return "refreshable" if cred.get("refresh_token") else "absent" + + async def get_user_oauth_credential( prisma_client: PrismaClient, user_id: str, @@ -1715,12 +1739,11 @@ async def resolve_valid_user_oauth_token( dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh actually happens, so the valid-token path never requires a DB handle. """ - if not cred or not cred.get("access_token"): + grant: Final = oauth_grant_state(cred) + if cred is None or grant == "absent": return None - if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): + if grant == "valid": return cred - if not cred.get("refresh_token"): - return None if prisma_client is None: from litellm.proxy.utils import get_prisma_client_or_throw @@ -1894,9 +1917,7 @@ async def get_mcp_submissions( order={"submitted_at": "desc"}, take=500, # safety cap; paginate if needed in a future iteration ) - items: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in rows] - for item in items: - decrypt_global_env_var_values(item.env_vars) + items: Final = list(_readable_mcp_servers(rows)) pending: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.pending_review) active: Final = sum(1 for i in items if i.approval_status == MCPApprovalStatus.active) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index e10bfd41ed6..cab4b6c161a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -43,9 +43,11 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + VendorCredentialState, aggregate_authorize, aggregate_token, complete_connect_flow, + describe_connect_flow, introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, @@ -798,21 +800,7 @@ def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MC return RedirectResponse(_append_query_params(redirect_uri, params), status_code=302) -async def _bridge_authorize_access_denial( - litellm_user_id: str, - mcp_server: MCPServer, - redirect_uri: str, - state: str, -) -> RedirectResponse | None: - """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed. - - Admits the user exactly as MCP egress will (the same ``reload_admitted_user`` constructor and the - same ``get_allowed_mcp_servers`` resolver), so an envelope is minted only when the resulting - session can actually list and call the server's tools. Without this gate the flow completes, the - client shows connected, and every tool request fail-closes to an empty list with nothing telling - the operator why. An availability fault (5xx, e.g. a DB outage's 503) propagates; an unknown or - deactivated user denies like a missing grant, fail closed. - """ +async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool: from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) @@ -821,13 +809,22 @@ async def _bridge_authorize_access_denial( ) try: - admitted: Final = await MCPRequestHandler.reload_admitted_user(litellm_user_id) + admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id) except HTTPException as exc: if exc.status_code >= 500: raise - return _bridge_access_denied_redirect(redirect_uri, state, mcp_server) - allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(admitted) - if mcp_server.server_id in allowed_server_ids: + return False + return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted) + + +async def _bridge_authorize_access_denial( + litellm_user_id: str, + mcp_server: MCPServer, + redirect_uri: str, + state: str, +) -> RedirectResponse | None: + """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed.""" + if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id): return None return _bridge_access_denied_redirect(redirect_uri, state, mcp_server) @@ -1910,6 +1907,38 @@ async def token_endpoint( ) +async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState: + """Whether the gateway itself can see a live vendor credential for this user and server. + + The one reading of "authorized" the connect page displays and the finish step enforces, so + the button a user sees and the grant they get cannot disagree. A read fault is neither, and + fails the scoped grant closed.""" + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # circular import at module load + get_user_oauth_credential, + oauth_grant_state, + ) + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # circular import at module load + + if prisma_client is None: + return "unavailable" + try: + credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id) + except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed + return "unavailable" + return "absent" if oauth_grant_state(credential) == "absent" else "present" + + +@router.get("/authorize/flow") +async def authorize_flow(request: Request, flow: str) -> Response: + return await describe_connect_flow( + request=request, + flow_handle=flow, + session_user_id=_session_cookie_user_id(request), + lookup_vendor_credential=_vendor_credential_state, + lookup_server_reachability=_user_can_reach_mcp_server, + ) + + @router.post("/authorize/complete") async def authorize_complete( request: Request, @@ -1934,6 +1963,8 @@ async def authorize_complete( delivery=delivery, team_id=team_id, decision=decision, + lookup_vendor_credential=_vendor_credential_state, + lookup_server_reachability=_user_can_reach_mcp_server, ) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index c7b0045dde5..3d94fa345d0 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -152,10 +152,8 @@ _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code" ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"] ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]] -"""Injected live-user revalidation (the token endpoint's mirror of admission): -``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is -a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else -fails the grant closed.""" +VendorCredentialState = Literal["present", "absent", "unavailable"] +"""The per-user vendor credential read has three outcomes: present, absent, or unavailable.""" _DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry" _DB_FAULTED_DESCRIPTION: Final = ( @@ -195,6 +193,16 @@ class ConsentTeam(BaseModel): team_alias: str | None = None +class LookupVendorCredential(Protocol): + """Injected read of a user's vendor credential for one server.""" + + def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ... + + +class LookupServerReachability(Protocol): + def __call__(self, user_id: str, server_id: str, /) -> Awaitable[bool]: ... + + class LookupConsentTeams(Protocol): """Injected lookup of the teams a signed-in user may bind a proxy-API credential to.""" @@ -205,6 +213,14 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr return "unresolvable" +async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: + return "unavailable" + + +async def _unreachable_server(user_id: str, server_id: str) -> bool: + return False + + class GatewayDcrClient(BaseModel): """The registration record sealed into a gateway DCR ``client_id``. @@ -449,7 +465,10 @@ def aggregate_authorize( A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the flow to that one server: the scope is sealed into the flow, carried into the code, and - bound into the session token, while the connect page interlude runs exactly as before. + bound into the session token. The connect URL carries only the flow handle; the page + learns the client origin, the scoped server, and whether its vendor OAuth is done from + :func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry + steers which server the page authorizes or names on the confirmation. Validation failures respond directly with 400 and never redirect: per RFC 6749 section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and @@ -474,10 +493,7 @@ def aggregate_authorize( resource_server_id=scoped_server.server_id if scoped_server is not None else None, audience=None, ) - connect_url: Final = _append_query_params( - f"{base_url}/ui/connect", - (("connect_flow", handle), ("connect_client", _origin_only(redirect_uri))), - ) + connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) response: Final = RedirectResponse(connect_url, status_code=303) _set_flow_cookie(response, request, handle, flow) return response @@ -684,6 +700,99 @@ def _origin_only(url: str) -> str: return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else "" +def _open_flow_for( + request: Request, flow_handle: str, session_user_id: str | None, now: datetime +) -> _ConnectFlow | Response: + sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle)) + if sealed_flow is None: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) + if flow is None or now.timestamp() >= flow.exp: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + if session_user_id is None: + return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") + if session_user_id != flow.user_id: + return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + return flow + + +async def _flow_target( + flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability +) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]: + if flow.resource_server_id is None: + return "unscoped", None + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle + MCPServerManager, + global_mcp_server_manager, + ) + + server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id) + if ( + server is None + or not server.is_gateway_managed_oauth2 + or not await lookup_server_reachability(flow.user_id, server.server_id) + ): + return "stale", None + state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive" + return state, server + + +class ConnectFlowDescription(TypedDict): + """What the connect page is allowed to know about one in-flight flow.""" + + state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]] + client_origin: ReadOnly[str] + server_id: ReadOnly[str | None] + server_name: ReadOnly[str | None] + connected: ReadOnly[bool | None] + + +async def _describe_opened_flow( + flow: _ConnectFlow, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> ConnectFlowDescription | Response: + state, server = await _flow_target(flow, lookup_server_reachability) + if state == "interactive" and server is not None: + credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id) + if credential == "unavailable": + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + interactive_description: Final[ConnectFlowDescription] = { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": server.server_id, + "server_name": server.server_name or server.alias or server.name, + "connected": credential == "present", + } + return interactive_description + described: Final[ConnectFlowDescription] = { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": None if server is None else server.server_id, + "server_name": None if server is None else (server.server_name or server.alias or server.name), + "connected": state == "m2m" or None, + } + return described + + +async def describe_connect_flow( + request: Request, + flow_handle: str, + session_user_id: str | None, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> Response: + opened: Final = _open_flow_for(request, flow_handle, session_user_id, datetime.now(timezone.utc)) + if isinstance(opened, Response): + return opened + described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability) + return ( + described + if isinstance(described, Response) + else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS) + ) + + async def complete_connect_flow( request: Request, flow_handle: str, @@ -692,56 +801,34 @@ async def complete_connect_flow( delivery: str | None = None, team_id: str | None = None, decision: str | None = None, + lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential, + lookup_server_reachability: LookupServerReachability = _unreachable_server, ) -> Response: - """The deliberate finish step of the connect flow: mint the gateway authorization - code and send the browser back to the client. + """Mint the code only after a deliberate POST by the sealed user. - Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly - per-flow cookie plus an exact match between the signed-in user and the user sealed - into the flow: a link crafted by another party dies here with ``access_denied`` - instead of minting a code for the victim's identity. The flow is single-use (an atomic - claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in. - - ``delivery`` chooses how the code reaches the client. Default (absent or - ``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"`` - renders the callback URL on a page instead, for a client whose redirect URI is a - loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box, - container): the 303 would dereference the browser machine's loopback and the code - would never arrive, so the user carries it over by pasting the URL into the client or - fetching it from the client machine's terminal. Manual delivery is honored only for - loopback redirect URIs; a routable redirect URI works from any browser by - construction, so those flows always redirect. The user who sees the page is exactly - the user the 303 would have carried the code to, and the same user already sees the - code today in the dead redirect's address bar, so the page exposes the code to no new - party. Unknown ``delivery`` values are rejected rather than defaulted: a client that - asked for manual delivery and got a dead redirect instead would silently lose its - code. - - ``decision`` and ``team_id`` come from the native-client consent page. ``"deny"`` - burns the flow and sends the client ``error=access_denied`` so it stops waiting; - ``team_id`` is sealed into the code only for proxy-API flows, where it picks which of - the user's teams the minted credential is attributed to. + A scoped flow additionally requires its sealed server to have a live vendor credential + before a code can be minted. The check happens before the single-use claim, so a + premature submit can be retried after authorization; denial deliberately bypasses it. """ if delivery not in (None, "redirect", "manual"): return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'") if decision not in (None, "approve", "deny"): return _oauth_error(400, "invalid_request", "decision must be 'approve' or 'deny'") - sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle)) - if sealed_flow is None: - return _oauth_error(400, "invalid_request", "unknown or expired connect flow") - flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) - if flow is None: - return _oauth_error(400, "invalid_request", "unknown or expired connect flow") now: Final = datetime.now(timezone.utc) - if now.timestamp() >= flow.exp: - return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection") - if session_user_id is None: - return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") - if session_user_id != flow.user_id: - return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + opened: Final = _open_flow_for(request, flow_handle, session_user_id, now) + if isinstance(opened, Response): + return opened + if decision != "deny": + described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability) + if isinstance(described, Response): + return described + if described["state"] == "stale": + return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available") + if described["connected"] is False: + return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing") flow_refusal: Final = _claim_refusal( await _SingleUseGuard(cache).claim( - f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS ), replayed=_oauth_error( 400, "invalid_request", "this connect flow was already completed; restart the connection" @@ -750,7 +837,7 @@ async def complete_connect_flow( if flow_refusal is not None: return flow_refusal response: Final = ( - _denied_flow_response(flow) if decision == "deny" else _approved_flow_response(flow, delivery, team_id, now) + _denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now) ) path, secure = _cookie_path_and_secure(request) response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax") diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index e7fd650a324..d7f238f142c 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -6272,8 +6272,7 @@ class MCPServerManager: ] } ) - db_mcp_servers: Final = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in raw_rows] - verbose_logger.info("Found %s MCP servers in database", len(db_mcp_servers)) + verbose_logger.info("Found %s MCP servers in database", len(raw_rows)) previous_registry: Final = self.registry new_registry: Final[dict[str, MCPServer]] = {} @@ -6281,8 +6280,9 @@ class MCPServerManager: # Stage one: build every server. Stage two assigns short prefixes # against the *full* set so dedup is deterministic regardless of # iteration order. - for server in db_mcp_servers: + for row in raw_rows: try: + server = LiteLLM_MCPServerTable.model_validate(row.model_dump()) existing_server = previous_registry.get(server.server_id) if ( @@ -6320,8 +6320,8 @@ class MCPServerManager: except Exception as e: verbose_logger.exception( "Skipping MCP server %s (%s) during DB reload: %s", - server.server_id, - getattr(server, "alias", None), + getattr(row, "server_id", None), + getattr(row, "alias", None), e, ) 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 60a9af89cc3..a44a041fdff 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -23,6 +23,7 @@ from pydantic import AnyUrl, ConfigDict from starlette.requests import Request as StarletteRequest from starlette.responses import JSONResponse from starlette.types import Message, Receive, Scope, Send +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG @@ -816,6 +817,11 @@ if MCP_AVAILABLE: } } return ListToolsResult.model_validate({"tools": listing.tools, "_meta": outcome_meta}) + except HTTPException as e: + from mcp.shared.exceptions import McpError + from mcp.types import INVALID_REQUEST, ErrorData + + raise McpError(ErrorData(code=INVALID_REQUEST, message=_http_detail_message(e.detail))) from e except Exception as e: verbose_logger.exception("Error in list_tools endpoint: %s", e) # Return empty list instead of failing completely @@ -911,15 +917,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 +970,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, @@ -1086,6 +1101,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers=mcp_server_auth_headers, oauth2_headers=oauth2_headers, raw_headers=raw_headers, + client_ip=_client_ip, host_progress_callback=host_progress_callback, **data, # for logging ) @@ -1119,7 +1135,7 @@ if MCP_AVAILABLE: except HTTPException as e: verbose_logger.error("HTTPException in MCP tool call: %s", e) return CallToolResult( - content=[TextContent(text=f"Error: {e.detail}", type="text")], + content=[TextContent(text=f"Error: {_http_detail_message(e.detail)}", type="text")], isError=True, ) except MCPUpstreamAuthError as e: @@ -1383,7 +1399,7 @@ if MCP_AVAILABLE: ######################################################## async def _get_allowed_mcp_servers_from_mcp_server_names( - mcp_servers: list[str] | None, + mcp_servers: Sequence[str] | None, allowed_mcp_servers: list[MCPServer], ) -> list[MCPServer]: """ @@ -1404,13 +1420,10 @@ if MCP_AVAILABLE: server_name_matched = False for server in allowed_mcp_servers: - if server: - match_list = [s.lower() for s in iter_known_server_prefixes(server) if s] - - if server_or_group.lower() in match_list: - filtered_server[server.server_id] = server - server_name_matched = True - break + if server and _server_answers_to(server, server_or_group): + filtered_server[server.server_id] = server + server_name_matched = True + break if not server_name_matched: try: @@ -1440,6 +1453,72 @@ if MCP_AVAILABLE: return allowed_mcp_servers + def _http_detail_message(detail: object) -> str: + return str(detail.get("error")) if isinstance(detail, dict) and detail.get("error") else str(detail) + + def _server_answers_to(server: MCPServer, name: str) -> bool: + requested: Final = name.lower() + return any(requested == known.lower() for known in iter_known_server_prefixes(server) if known) + + class _McpDeniedDetail(TypedDict): + error: ReadOnly[str] + + async def raise_denied_scoped_mcp_access( + requested_names: Sequence[str], + user_api_key_auth: UserAPIKeyAuth | None, + client_ip: str | None = None, + ) -> None: + """A scoped request (``/mcp/`` path or ``x-mcp-servers`` header) resolved to zero + allowed servers, so the denial must be loud: a silent 200 with no tools reads as a healthy + server with no tools. Unknown, unauthorized, and access-group names all share one generic + error so scoping cannot probe which servers exist; the agent variant fires only when the + same request resolves once the agent binding is stripped, proving the binding caused the veto.""" + agent_id: Final = user_api_key_auth.agent_id if user_api_key_auth else None + if user_api_key_auth is not None and agent_id: + resolved_without_agent: Final = await _get_allowed_mcp_servers( + user_api_key_auth=user_api_key_auth.model_copy(update=types.MappingProxyType({"agent_id": None})), + mcp_servers=requested_names, + client_ip=client_ip, + ) + + def _resolved_to_server(name: str) -> bool: + return any(_server_answers_to(server, name) for server in resolved_without_agent) + + vetoed_server: Final = next((name for name in requested_names if _resolved_to_server(name)), None) + if vetoed_server is not None: + agent_denial: Final[_McpDeniedDetail] = { + "error": ( + f"MCP server '{vetoed_server}' is not available to this key: the key is bound to " + f"agent '{agent_id}', whose MCP grants do not include this server. Add the server " + f"to the agent's object_permission.mcp_servers (edit the agent in the Admin UI or " + f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." + ) + } + raise HTTPException(status_code=403, detail=agent_denial) + vetoed_group: Final = next( + ( + name + for name in requested_names + if not _resolved_to_server(name) + and any(name in (server.access_groups or ()) for server in resolved_without_agent) + ), + None, + ) + if vetoed_group is not None: + group_denial: Final[_McpDeniedDetail] = { + "error": ( + f"MCP access group '{vetoed_group}' is not available to this key: the key is bound to " + f"agent '{agent_id}', whose MCP grants do not include it. Add the group to the " + f"agent's object_permission.mcp_access_groups (edit the agent in the Admin UI or " + f"PATCH /v1/agents/{agent_id}), or use a key that is not bound to the agent." + ) + } + raise HTTPException(status_code=403, detail=group_denial) + generic_denial: Final[_McpDeniedDetail] = { + "error": f"The key is not allowed to access the requested MCP servers: {', '.join(requested_names)}" + } + raise HTTPException(status_code=403, detail=generic_denial) + def _tool_name_matches(tool_name: str, filter_list: list[str], mcp_server: MCPServer) -> bool: """ Check if a tool name matches any name in the filter list. @@ -1532,7 +1611,7 @@ if MCP_AVAILABLE: async def _get_allowed_mcp_servers( user_api_key_auth: UserAPIKeyAuth | None, - mcp_servers: list[str] | None, + mcp_servers: Sequence[str] | None, client_ip: str | None = None, ) -> list[MCPServer]: """Return allowed MCP servers for a request after applying filters. @@ -1968,6 +2047,12 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, client_ip=client_ip, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) # Pre-fetch OAuth credentials only when at least one server uses OAuth2, # to avoid an unnecessary DB round-trip on requests with no OAuth2 MCP servers. @@ -2395,6 +2480,8 @@ if MCP_AVAILABLE: ) verbose_logger.debug("Successfully fetched %s tools from managed MCP servers", len(listing.tools)) return listing + except HTTPException: + raise except Exception as e: verbose_logger.exception("Error getting tools from managed MCP servers: %s", e) # Continue with an empty listing instead of failing completely @@ -3077,6 +3164,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, + client_ip: str | None = None, **kwargs: Any, ) -> CallToolResult: """ @@ -3107,6 +3195,12 @@ if MCP_AVAILABLE: mcp_servers=mcp_servers, allowed_mcp_servers=allowed_mcp_servers, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_auth, + client_ip=client_ip, + ) if not allowed_mcp_servers: raise HTTPException( status_code=403, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index af02c11ad86..24e2f5ce64d 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, ) @@ -308,6 +366,7 @@ async def handle_mcp_tool_call( from litellm.proxy._experimental.mcp_server.server import ( _get_allowed_mcp_servers, execute_mcp_tool, + raise_denied_scoped_mcp_access, ) allowed_mcp_servers: Final = await _get_allowed_mcp_servers( @@ -315,6 +374,12 @@ async def handle_mcp_tool_call( mcp_servers=mcp_servers, client_ip=client_ip, ) + if mcp_servers and not allowed_mcp_servers: + await raise_denied_scoped_mcp_access( + requested_names=mcp_servers, + user_api_key_auth=user_api_key_dict, + client_ip=client_ip, + ) # Reject before dispatch when the key has no accessible servers; otherwise an # unprefixed local tool name would fall through to the local registry in diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 91a97ad6544..71475320c2c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2634,6 +2634,20 @@ ], "title": "Mcp Tool Permissions" }, + "mcp_toolsets": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Mcp Toolsets" + }, "models": { "anyOf": [ { @@ -4687,6 +4701,17 @@ "title": "Created At", "type": "string" }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, "display_title": { "anyOf": [ { @@ -4713,6 +4738,17 @@ ], "title": "Latest Version" }, + "search_score": { + "anyOf": [ + { + "type": "number" + }, + { + "type": "null" + } + ], + "title": "Search Score" + }, "source": { "title": "Source", "type": "string" @@ -4781,7 +4817,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 +4885,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 +9666,9 @@ { "enum": [ "warn", - "end_session" + "end_session", + "block", + "alert" ], "type": "string" }, @@ -9605,7 +9676,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 +11985,9 @@ { "enum": [ "warn", - "end_session" + "end_session", + "block", + "alert" ], "type": "string" }, @@ -11922,7 +11995,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": { @@ -19813,6 +19886,46 @@ ] } }, + "/authorize/flow": { + "get": { + "operationId": "authorize_flow_authorize_flow_get", + "parameters": [ + { + "in": "query", + "name": "flow", + "required": true, + "schema": { + "title": "Flow", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Authorize Flow", + "tags": [ + "mcp_discoverable" + ] + } + }, "/callback": { "get": { "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c28ac8848ba..d1b391d57a3 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 @@ -2831,7 +2833,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "Enable only if your deployment is experiencing phantom " "BudgetExceededError responses caused by leaked reservations " "(see GitHub issue #27639). " - "A proxy-level WARNING is logged on every request while this flag " + "An INFO notice is logged once per worker at config load while this flag " "is active as a reminder that hard enforcement is relaxed." ), ) @@ -3583,6 +3585,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase): litellm_callback_params=[ "OTEL_EXPORTER", "OTEL_ENDPOINT", + "OTEL_TRACES_ENDPOINT", "OTEL_HEADERS", ], ) 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..f78c4221f5a 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -11,7 +11,7 @@ from fastapi import HTTPException, Request, status from pydantic import PositiveInt, TypeAdapter, ValidationError import litellm -from litellm import Router, provider_list +from litellm import Router, constants, provider_list from litellm._logging import verbose_proxy_logger from litellm.constants import ( BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, @@ -1390,6 +1390,24 @@ def warn_once_if_custom_auth_skips_common_checks( _custom_auth_common_checks_warning_emitted = True +def log_once_if_budget_reservation_disabled( + *, + disabled: bool, + logger: Logger = verbose_proxy_logger, +) -> None: + if constants.budget_reservation_disabled_info_emitted or not disabled: + return + logger.info( + "disable_budget_reservation is enabled: skipping optimistic budget " + "reservation. Budget enforcement is read-time only. Concurrent " + "requests can each pass the spend check before their cost is recorded, " + "so a configured budget may be briefly exceeded under high concurrency. " + "Set disable_budget_reservation to False or remove it to restore " + "hard per-request budget enforcement." + ) + constants.budget_reservation_disabled_info_emitted = True # rebind-ok: process-wide one-shot sentinel + + def is_pass_through_provider_route(route: str) -> bool: PROVIDER_SPECIFIC_PASS_THROUGH_ROUTES: Final = [ "vertex-ai", @@ -1981,9 +1999,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/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 93293db24c6..b39b1f330b3 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2706,14 +2706,6 @@ async def _reserve_budget_after_common_checks( if skip_budget_checks: return if general_settings.get("disable_budget_reservation") is True: - verbose_proxy_logger.warning( - "disable_budget_reservation is enabled: skipping optimistic budget " - "reservation. Budget enforcement is read-time only — concurrent " - "requests can each pass the spend check before their cost is recorded, " - "so a configured budget may be briefly exceeded under high concurrency. " - "Set disable_budget_reservation to False or remove it to restore " - "hard per-request budget enforcement." - ) return from litellm.proxy.spend_tracking.budget_reservation import ( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5e6c9b34332..d0e3914e9fb 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(): @@ -3460,9 +3508,13 @@ class ProxyBaseLLMRequestProcessing: error_body: Final = await http_status_error.response.aread() error_text: Final = error_body.decode("utf-8") + error_headers: Final = { # mutable-ok: HTTPException takes a plain header dict + k: v if isinstance(v, str) else str(v) for k, v in safe_headers.items() + } raise HTTPException( status_code=http_status_error.response.status_code, detail={"error": error_text}, + headers=error_headers, ) error_msg: Final = f"{e}" # Check for AttributeError in the exception chain. diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index 39e74d2c8bd..d9d56778405 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -10,6 +10,7 @@ import litellm from litellm import get_secret from litellm._logging import verbose_proxy_logger from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, @@ -507,6 +508,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset( PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, + CLIENT_OUTPUT_CEILING_METADATA_KEY, "disable_global_guardrails", "disable_global_guardrail", "opted_out_global_guardrails", diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 836a4a778bb..fd9b3beee46 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -1,7 +1,10 @@ import base64 import os +from collections.abc import Mapping from typing import Final, Literal, cast +from pydantic import TypeAdapter, ValidationError + from litellm._logging import verbose_proxy_logger # Versioned ciphertext marker for AES-256-GCM values. @@ -203,3 +206,40 @@ def decrypt_value(value: bytes, signing_key: str) -> str: return plaintext except Exception as e: raise e + + +class SecretMapDecodeError(RuntimeError): + pass + + +_SECRET_MAP: Final = TypeAdapter(Mapping[str, str]) +_STORED_SECRET_MAP: Final = TypeAdapter(Mapping[str, str] | str) +_SECRET_STRING: Final = TypeAdapter(str) + + +def encrypt_secret_map(value: Mapping[str, str], new_encryption_key: str | None = None) -> str: + if not value: + return "{}" + ciphertext: Final = _SECRET_STRING.validate_python( + encrypt_value_helper(_SECRET_MAP.dump_json(value).decode(), new_encryption_key=new_encryption_key), strict=True + ) + return _SECRET_STRING.dump_json(ciphertext).decode() + + +def decode_secret_map(value: object, *, key: str) -> Mapping[str, str] | None: + if value is None: + return None + try: + stored: Final = ( + _STORED_SECRET_MAP.validate_json(value, strict=True) + if isinstance(value, str) and value.lstrip().startswith(("{", '"')) + else _STORED_SECRET_MAP.validate_python(value, strict=True) + ) + if not isinstance(stored, str): + return stored + decrypted: Final = decrypt_value_helper( + value=stored, key=key, exception_type="debug", return_original_value=False + ) + return _SECRET_MAP.validate_json(decrypted, strict=True) + except ValidationError: + raise SecretMapDecodeError(f"Cannot decode encrypted MCP {key}; check LITELLM_SALT_KEY") from None diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 460b348e188..63da5d15207 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -125,6 +125,7 @@ async def _resync_model_deployments(model_name: str) -> bool: ) return proxy_server.llm_router is not None async with proxy_server.MODEL_RECONCILE_LOCK: + await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client) proxy_server.proxy_config._add_deployment(db_models=rows) proxy_server.llm_model_list = router.get_model_list() return True 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..b866ecc741f 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 @@ -73,6 +75,8 @@ SELECT COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens, COALESCE(SUM(spend), 0)::float8 AS spend, COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, + COALESCE(SUM(classifier_cost), 0)::float8 AS classifier_cost, + COALESCE(SUM(classifier_cost_recorded_turns), 0)::int AS classifier_cost_recorded_turns, COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds FROM windowed GROUP BY router_name, router_type @@ -93,6 +97,7 @@ class AutoRouterTurnTransaction: total_tokens: int spend: float saved_spend: float + classifier_cost: float covered: bool cache_hit: bool cache_ttl_seconds: int | None @@ -223,6 +228,7 @@ def build_autorouter_turn_transaction( total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0), spend=float(payload.get("spend") or 0.0) + (classifier_cost or 0.0), saved_spend=saved_spend, + classifier_cost=classifier_cost or 0.0, covered=cache.covered, cache_hit=cache.read_tokens > 0, cache_ttl_seconds=cache.write_ttl_seconds, @@ -264,7 +270,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t ( last_model, models, turns, unordered_turns, covered_turns, cache_hits, same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, return_turns, return_hits, return_expired_misses, return_within_ttl_misses, - ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns + ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns ) VALUES ( {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, @@ -275,13 +281,15 @@ VALUES ( (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END), (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END), {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8, - {_TIER_DELTA} + {_p("classifier_cost")}::float8, 1, {_TIER_DELTA} ) ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, total_tokens = t.total_tokens + EXCLUDED.total_tokens, spend = t.spend + EXCLUDED.spend, saved_spend = t.saved_spend + EXCLUDED.saved_spend, + classifier_cost = t.classifier_cost + EXCLUDED.classifier_cost, + classifier_cost_recorded_turns = t.classifier_cost_recorded_turns + 1, covered_turns = t.covered_turns + EXCLUDED.covered_turns, cache_hits = t.cache_hits + EXCLUDED.cache_hits, ttl_5m_turns = t.ttl_5m_turns + EXCLUDED.ttl_5m_turns, diff --git a/litellm/proxy/db/check_migration.py b/litellm/proxy/db/check_migration.py index b7a07d4eeea..6e2a06c96e1 100644 --- a/litellm/proxy/db/check_migration.py +++ b/litellm/proxy/db/check_migration.py @@ -46,17 +46,32 @@ def extract_sql_commands(diff_output: str) -> list[str]: def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: """Checks for differences between current database and Prisma schema. + + Never raises: a diff that cannot be produced, because the runner is missing, + because the command failed, or because it outlived its budget, is reported as + "no diff" so boot continues. + Returns: A tuple containing: - A boolean indicating if differences were found (True) or not (False). - - A string with the diff output or error message. - Raises: - subprocess.CalledProcessError: If the Prisma command fails. - Exception: For any other errors during execution. + - The SQL commands that would close the diff, empty when there is none. """ - verbose_logger.debug("Checking for Prisma schema diff...") try: - result: Final = subprocess.run( + from litellm_proxy_extras.prisma_toolchain import ( + PRISMA_COMMAND_TIMEOUT_ENV_VAR, + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Skipping the migration diff: litellm-proxy-extras has no Prisma runner. Error: {e}" + ) + return False, [] + + verbose_logger.debug("Checking for Prisma schema diff...") + timeout: Final = prisma_command_timeout() + try: + result: Final = run_prisma( [ "prisma", "migrate", @@ -67,12 +82,10 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: "./schema.prisma", "--script", ], - capture_output=True, - text=True, - check=True, + timeout=timeout, + env=os.environ.copy(), ) - # return True, "Migration diff generated successfully." sql_commands: Final = extract_sql_commands(result.stdout) if sql_commands: @@ -83,6 +96,12 @@ def check_prisma_schema_diff_helper(db_url: str) -> tuple[bool, list[str]]: return True, sql_commands else: return False, [] + except subprocess.TimeoutExpired: + print( # noqa: T201 # boot-time operator output, same channel as this helper's other messages + f"Timed out after {timeout}s generating the migration diff. " + f"Raise {PRISMA_COMMAND_TIMEOUT_ENV_VAR} if this database needs longer." + ) + return False, [] except subprocess.CalledProcessError as e: error_message: Final = f"Failed to generate migration diff. Error: {e.stderr}" print(error_message) # noqa: T201 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 9230be8055e..914c961b145 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -502,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, @@ -2188,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/prisma_client.py b/litellm/proxy/db/prisma_client.py index 2190ae55fd2..9ea2432f2f5 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -937,9 +937,17 @@ class PrismaManager: use_v2_resolver=use_v2_resolver, ) else: + try: + from litellm_proxy_extras.prisma_toolchain import ( + prisma_command_timeout, + run_prisma, + ) + except ImportError as e: + verbose_proxy_logger.error("\x1b[1;31mLiteLLM: Failed to import proxy extras. Got %s\x1b[0m", e) + return False + PrismaManager._raise_if_partitioned_spend_logs() - # Use prisma db push with increased timeout - subprocess.run( + run_prisma( [ "prisma", "db", @@ -947,13 +955,15 @@ class PrismaManager: "--accept-data-loss", "--skip-generate", ], - timeout=60, - check=True, + timeout=prisma_command_timeout(), + env=os.environ.copy(), + stdout=None, + stderr=None, ) PrismaManager._apply_replica_identity_full_if_requested() return True - except subprocess.TimeoutExpired: - verbose_proxy_logger.warning("Attempt %s timed out", attempt + 1) + except subprocess.TimeoutExpired as e: + verbose_proxy_logger.warning("Attempt %s timed out after %.0fs", attempt + 1, e.timeout) time.sleep(random.randrange(5, 15)) except subprocess.CalledProcessError as e: attempts_left = 3 - attempt 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/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index a8b33109900..e20f0b320b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -547,7 +547,7 @@ class ToolPermissionGuardrail(CustomGuardrail): for _tool_call, is_allowed, _rule_id, message in checked: if not is_allowed and message is not None: - verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": raise GuardrailRaisedException( guardrail_name=self.guardrail_name, message=message, blocked_content=True @@ -809,7 +809,7 @@ class ToolPermissionGuardrail(CustomGuardrail): new_tools: Final = self._collect_request_tools(data) if not new_tools: - verbose_proxy_logger.warning( + verbose_proxy_logger.debug( "Tool Permission Guardrail: not running guardrail. No tools or functions in data" ) return data @@ -820,7 +820,7 @@ class ToolPermissionGuardrail(CustomGuardrail): is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) if not is_allowed and message is not None: - verbose_proxy_logger.warning("Tool Permission Guardrail: %s", message) + verbose_proxy_logger.info("Tool Permission Guardrail: %s", message) if self.on_disallowed_action == "block": raise HTTPException( status_code=400, 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/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 8b82842353c..bec12ba8201 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,7 +5,7 @@ Pre-call hook that filters MCP tools semantically before LLM inference. Reduces context window size and improves tool selection accuracy. """ -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Final, Optional from fastapi import HTTPException @@ -104,7 +104,7 @@ class SemanticToolFilterHook(CustomLogger): ) # Parse to separate MCP tools from other tools - mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + mcp_tools, _ = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_tools: return [] @@ -173,7 +173,11 @@ class SemanticToolFilterHook(CustomLogger): return [name for name in names if name] @staticmethod - def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]: + async def _narrow_mcp_references( + tools: Sequence[Mapping[str, object]], + selected_tool_names: list[str], + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] | None = None, + ) -> list[object]: """ Restrict each litellm_proxy MCP reference to the semantically selected tools. @@ -192,13 +196,14 @@ class SemanticToolFilterHook(CustomLogger): LiteLLM_Proxy_MCP_Handler, ) + via_gateway: Final = await ( + LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools, served_names) + if served_names is not None + else LiteLLM_Proxy_MCP_Handler.routes_through_gateway(tools) + ) return [ - ( - {**tool, "allowed_tools": selected_tool_names} - if isinstance(tool, dict) and LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([tool]) - else tool - ) - for tool in tools + {**tool, "allowed_tools": selected_tool_names} if isinstance(tool, dict) and routed else tool + for tool, routed in zip(tools, via_gateway, strict=True) ] def _is_mcp_tool(self, tool: object) -> bool: @@ -325,7 +330,7 @@ class SemanticToolFilterHook(CustomLogger): filtered_expanded_tools = await self._filter_expanded_tools(data=data, expanded_tools=expanded_tools) selected_tool_names: Final = self._selected_tool_names(filtered_expanded_tools) - narrowed_tools: Final = self._narrow_mcp_references(tools, selected_tool_names) + narrowed_tools: Final = await self._narrow_mcp_references(tools, selected_tool_names) data["tools"] = narrowed_tools self._emit_filter_metadata_safe( data=data, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 56512570448..893174dab03 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -18,6 +18,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm._uuid import uuid from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, LITELLM_PROXY_MASTER_KEY_ALIAS, @@ -325,7 +326,9 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg # ``attempted_fallbacks`` and ``original_model_group`` are written by the router # and read by spend logs as fact; a client value has no legitimate meaning and no # key or team setting keeps it, so the strip is never gated. -_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset({"attempted_fallbacks", "original_model_group"}) +_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset( + {"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY} +) _ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override" # Request fields whose value, when URL-valued, becomes the outbound destination @@ -789,12 +792,6 @@ def apply_missing_session_id_policy( ) -def is_claude_code_user_agent(user_agent: str) -> bool: - """Claude Code identifies itself as ``claude-cli/ ...``; the IDE - extensions and the Agent SDK run through the same CLI and share that prefix.""" - return user_agent.startswith("claude-cli/") - - def is_codex_user_agent(user_agent: str) -> bool: """Codex builds its user agent as ``/ ...`` and ships several first-party originators: ``codex-tui``, ``codex_cli_rs``, @@ -811,6 +808,8 @@ def should_auto_drop_params_for_agentic_cli(user_agent: str, data: dict, proxy_c requests routed to providers that reject them. An explicit drop_params from the caller or in the operator's ``litellm_settings`` always wins over this default.""" + from litellm.llms.anthropic.common_utils import is_claude_code_user_agent + if not (is_claude_code_user_agent(user_agent) or is_codex_user_agent(user_agent)): return False if "drop_params" in data: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 2a7813bc140..bbc914a772a 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -484,6 +484,8 @@ class _SessionAggRow(BaseModel): total_tokens: int spend: float saved_spend: float + classifier_cost: float + classifier_cost_recorded_turns: int session_seconds: float @@ -520,6 +522,7 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals: avg_tokens_per_session=row.total_tokens / sessions if sessions else 0.0, spend=row.spend, saved_spend=row.saved_spend, + classifier_cost=row.classifier_cost if row.classifier_cost_recorded_turns == row.turns else None, baseline_spend=baseline_spend, saved_pct=_pct(row.saved_spend, baseline_spend), saved_per_session=row.saved_spend / sessions if sessions else 0.0, @@ -552,6 +555,7 @@ def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup: avg_tokens_per_session=totals.avg_tokens_per_session, spend=totals.spend, saved_spend=totals.saved_spend, + classifier_cost=totals.classifier_cost, baseline_spend=totals.baseline_spend, saved_pct=totals.saved_pct, saved_per_session=totals.saved_per_session, @@ -582,6 +586,8 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow: total_tokens=sum(row.total_tokens for row in rows), spend=sum(row.spend for row in rows), saved_spend=sum(row.saved_spend for row in rows), + classifier_cost=sum(row.classifier_cost for row in rows), + classifier_cost_recorded_turns=sum(row.classifier_cost_recorded_turns for row in rows), session_seconds=sum(row.session_seconds for row in rows), ) @@ -645,6 +651,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 +688,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/credential_migration.py b/litellm/proxy/management_endpoints/credential_migration.py index e0725119576..915cce87dbd 100644 --- a/litellm/proxy/management_endpoints/credential_migration.py +++ b/litellm/proxy/management_endpoints/credential_migration.py @@ -43,7 +43,9 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( _ALGO_AES_GCM, _ENCRYPTION_ALGORITHM_SETTING, _V2_GCM_PREFIX, + SecretMapDecodeError, _get_salt_key, + decode_secret_map, decrypt_value_helper, encrypt_value_helper, ) @@ -65,6 +67,20 @@ class LocationReport: # Used by --check (read-only classification): legacy: int = 0 # nacl ciphertext still awaiting migration + def count(self, classification: ValueClass | None) -> None: + if classification is None: + return + self.scanned += 1 + match classification: + case "migrated": + self.already_v2 += 1 + case "legacy": + self.legacy += 1 + case "undecryptable": + self.undecryptable += 1 + case _: + self.plaintext += 1 + def as_dict(self) -> dict[str, int]: return { "scanned": self.scanned, @@ -441,7 +457,7 @@ def _classify_callback_value(value: object) -> ValueClass: _COVERED_TABLE_SPECS: Final = [ ("model_table", "litellm_proxymodeltable", ("litellm_params",), ()), ("credentials", "litellm_credentialstable", ("credential_values",), ()), - ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars"), ()), + ("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()), ("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)), ("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)), ] @@ -472,14 +488,18 @@ def _classify_into_report(report: LocationReport, value: str) -> None: names, base URLs, …) do not decrypt and fall through to ``plaintext``, so over-scanning a column is harmless to the residual count. """ - report.scanned += 1 - cls: Final = classify_value(value, key="scan") - if cls == "migrated": - report.already_v2 += 1 - elif cls == "legacy": - report.legacy += 1 - else: # plaintext / not-a-string - report.plaintext += 1 + report.count(classify_value(value, key="scan")) + + +def _classify_secret_map(value: object, key: str) -> ValueClass | None: + try: + decoded: Final = decode_secret_map(value, key=key) + except SecretMapDecodeError: + return "undecryptable" + if not decoded: + return None + ciphertext: Final = json.loads(value) if isinstance(value, str) and value.lstrip().startswith('"') else value + return "migrated" if is_migrated(ciphertext) else "legacy" async def _scan_one_table( @@ -503,6 +523,9 @@ async def _scan_one_table( raw = getattr(row, col, None) if raw is None: continue + if db_attr == "litellm_mcpservertable" and col in ("static_headers", "env"): + report.count(_classify_secret_map(raw, col)) + continue if isinstance(raw, str): try: raw = json.loads(raw) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8e51e250319..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) @@ -2862,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 @@ -2918,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, @@ -2933,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, @@ -2989,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 @@ -6432,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/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 32da2658b99..2ea46b740a8 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. @@ -942,7 +946,14 @@ async def handle_bedrock_count_tokens( except BedrockError as e: # Convert BedrockError to HTTPException for FastAPI verbose_proxy_logger.error("BedrockError in handle_bedrock_count_tokens: %s", e) - raise HTTPException(status_code=e.status_code, detail={"error": e.message}) + from litellm.litellm_core_utils.llm_response_utils.get_headers import get_response_headers + + provider_headers: Final = getattr(getattr(e, "response", None), "headers", None) + raise HTTPException( + status_code=e.status_code, + detail={"error": e.message}, + headers=get_response_headers(provider_headers) if provider_headers else None, + ) except HTTPException: # Re-raise HTTP exceptions as-is raise @@ -986,8 +997,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/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4be0f556ed7..1bcda38657b 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final, Literal import litellm @@ -16,7 +16,11 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import independent_snapshot +from litellm.litellm_core_utils.core_helpers import ( + get_metadata_variable_name_from_kwargs, + get_or_create_metadata_bucket, + independent_snapshot, +) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -25,6 +29,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineStep, PipelineStepResult, ) +from litellm.types.utils import StandardLoggingGuardrailInformation try: from fastapi.exceptions import HTTPException @@ -118,6 +123,7 @@ class PipelineExecutor: return _allow_result(step_results=step_results, working_data=working_data, request_data=data) if action == "block": + _carry_working_guardrail_information(working_data=working_data, request_data=data) return PipelineExecutionResult( terminal_action="block", step_results=step_results, @@ -126,6 +132,7 @@ class PipelineExecutor: ) if action == "modify_response": + _carry_working_guardrail_information(working_data=working_data, request_data=data) return PipelineExecutionResult( terminal_action="modify_response", step_results=step_results, @@ -168,34 +175,33 @@ class PipelineExecutor: verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail) return ("error", None, f"Guardrail '{step.guardrail}' not found", None) + # Inject guardrail name into metadata so should_run_guardrail() allows it + if "metadata" not in data: + data["metadata"] = {} + data["metadata"]["guardrails"] = [step.guardrail] + + # A scan_raw_request step evaluates the pristine pre-pipeline + # snapshot instead of `data` (which earlier pass_data steps in + # this same pipeline may have already rewritten), same reason + # the normal sequential/parallel guardrail loops do this. + scans_raw_request: Final = callback.scan_raw_request + hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data + independent_snapshot(raw_request_snapshot) + if scans_raw_request and raw_request_snapshot is not None + else data + ) + if hook_input is not data: + hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] + snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input)) + + # Use unified_guardrail path if callback implements apply_guardrail + target: CustomLogger = callback + use_unified: Final = "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + if use_unified: + hook_input["guardrail_to_apply"] = callback + target = UnifiedLLMGuardrails() + try: - # Inject guardrail name into metadata so should_run_guardrail() allows it - if "metadata" not in data: - data["metadata"] = {} - data["metadata"]["guardrails"] = [step.guardrail] - - # A scan_raw_request step evaluates the pristine pre-pipeline - # snapshot instead of `data` (which earlier pass_data steps in - # this same pipeline may have already rewritten), same reason - # the normal sequential/parallel guardrail loops do this. - scans_raw_request: Final = callback.scan_raw_request - hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data - independent_snapshot(raw_request_snapshot) - if scans_raw_request and raw_request_snapshot is not None - else data - ) - if hook_input is not data: - hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] - - # Use unified_guardrail path if callback implements apply_guardrail - target: CustomLogger = callback - use_unified: Final = ( - "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks - ) - if use_unified: - hook_input["guardrail_to_apply"] = callback - target = UnifiedLLMGuardrails() - if mode == "pre_call": response = await target.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -233,6 +239,12 @@ class PipelineExecutor: else: verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e) return ("error", None, str(e), e) + finally: + if hook_input is not data: + _append_guardrail_information( + request_data=data, + entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:], + ) @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: @@ -283,6 +295,40 @@ def _restore_request_guardrails( return {**working_data, "metadata": stripped} # mutable-ok: request dict +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +def _recorded_guardrail_information(source: Mapping[str, object]) -> list[StandardLoggingGuardrailInformation]: + bucket: Final = source.get(get_metadata_variable_name_from_kwargs(source)) + recorded: Final = bucket.get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None + return recorded if isinstance(recorded, list) else [] + + +def _append_guardrail_information( + request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data + entries: Sequence[StandardLoggingGuardrailInformation], +) -> None: + if not entries: + return + _, request_bucket = get_or_create_metadata_bucket(request_data) + existing: Final = request_bucket.get(_GUARDRAIL_INFORMATION_KEY) + if isinstance(existing, list): + existing.extend(entries) + return + request_bucket[_GUARDRAIL_INFORMATION_KEY] = list(entries) + + +def _carry_working_guardrail_information( + working_data: Mapping[str, object], + request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data +) -> None: + recorded: Final = _recorded_guardrail_information(working_data) + existing: Final = _recorded_guardrail_information(request_data) + if recorded is existing: + return + _append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing]) + + def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: """ Map pipeline step outcome to the configured action. diff --git a/litellm/proxy/prometheus_cleanup.py b/litellm/proxy/prometheus_cleanup.py index d9827723887..c65aaeedfaa 100644 --- a/litellm/proxy/prometheus_cleanup.py +++ b/litellm/proxy/prometheus_cleanup.py @@ -8,10 +8,13 @@ from __future__ import annotations import glob import os +import re from typing import Final from litellm._logging import verbose_proxy_logger +_LIVE_GAUGE_PID: Final = re.compile(r"gauge_live[a-z]*_(\d+)\.db$") + def wipe_directory(directory: str) -> None: """Delete all .db files in the directory. Called once before workers fork.""" @@ -38,3 +41,35 @@ def mark_worker_exit(worker_pid: int) -> None: verbose_proxy_logger.info("Prometheus cleanup: marked worker %s as dead", worker_pid) except Exception as e: verbose_proxy_logger.warning("Failed to mark prometheus worker %s as dead: %s", worker_pid, e) + + +def _is_running(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def mark_dead_workers(directory: str) -> tuple[int, ...]: + """Drop the live-gauge files of workers that no longer exist and return their pids. + + Uvicorn's multi-worker supervisor has no exit hook, so a replacement worker calls this at startup; without it + a crashed worker's in-flight gauges stay in the aggregate forever. + """ + owners: Final = frozenset( + int(match.group(1)) + for match in map(_LIVE_GAUGE_PID.search, glob.glob(os.path.join(directory, "gauge_live*_*.db"))) + if match is not None + ) + dead: Final = tuple(sorted(pid for pid in owners if pid != os.getpid() and not _is_running(pid))) + if not dead: + return dead + from prometheus_client import multiprocess + + for pid in dead: + multiprocess.mark_process_dead(pid, path=directory) + verbose_proxy_logger.info("Prometheus cleanup: marked dead workers %s in %s", dead, directory) + return dead diff --git a/litellm/proxy/prometheus_metrics_server.py b/litellm/proxy/prometheus_metrics_server.py index 4a9651d62e1..01479f9dd41 100644 --- a/litellm/proxy/prometheus_metrics_server.py +++ b/litellm/proxy/prometheus_metrics_server.py @@ -30,6 +30,7 @@ from litellm.integrations.prometheus_metrics_endpoint import make_metrics_asgi_a from litellm.llms.custom_httpx.http_handler import HTTPHandler METRICS_PATH: Final = "/metrics" +HEALTH_PATH: Final = "/health" PID_HEADER: Final = "x-litellm-metrics-pid" _PARENT_POLL_INTERVAL_SECONDS: Final = 1.0 _STARTUP_TIMEOUT_SECONDS: Final = 30.0 @@ -77,6 +78,10 @@ def build_metrics_app(multiproc_dir: str) -> FastAPI: app: Final = FastAPI(title="LiteLLM Prometheus metrics", docs_url=None, redoc_url=None, openapi_url=None) app.mount(METRICS_PATH, _add_pid_header(make_metrics_asgi_app(registry))) + @app.get(HEALTH_PATH) + def health() -> dict[str, str]: + return {"status": "healthy", "multiproc_dir": multiproc_dir} + return app diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0915b8dd1b9..38c4c2e96c9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -306,6 +306,7 @@ from litellm.proxy.auth.auth_checks import ( from litellm.proxy.auth.auth_utils import ( check_response_size_is_safe, is_request_body_safe, + log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, ) from litellm.proxy.auth.fallback_model_access import router_fallback_access_check @@ -630,6 +631,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( router as pass_through_router, ) +from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.public_endpoints.public_v1 import router as public_v1_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router @@ -1058,6 +1060,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: init_verbose_loggers() + prometheus_multiproc_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + if prometheus_multiproc_dir: + mark_dead_workers(prometheus_multiproc_dir) + ## RUN WORKER STARTUP HOOKS (e.g., gflags initialization) ## _startup_hooks_env: Final = os.environ.get("LITELLM_WORKER_STARTUP_HOOKS", "") if _startup_hooks_env: @@ -1364,6 +1370,9 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await proxy_shutdown_event(worker_heartbeat=worker_heartbeat) + if prometheus_multiproc_dir: + mark_worker_exit(os.getpid()) + def _generate_stable_operation_id(route: "APIRoute") -> str: operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}") @@ -5653,6 +5662,10 @@ class ProxyConfig: run_common_checks=bool(general_settings.get("custom_auth_run_common_checks", False)), ) + log_once_if_budget_reservation_disabled( + disabled=general_settings.get("disable_budget_reservation") is True, + ) + custom_key_generate: Final = general_settings.get("custom_key_generate", None) if custom_key_generate is not None: user_custom_key_generate = get_instance_fn(value=custom_key_generate, config_file_path=config_file_path) @@ -7104,11 +7117,10 @@ class ProxyConfig: ], ) - # Only load models from DB if "models" is in supported_db_objects (or if supported_db_objects is not set) - if self._should_load_db_object(object_type="models"): - new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) - - # update llm router + load_models: Final = self._should_load_db_object(object_type="models") + new_models: Final = await self._get_models_from_db(prisma_client=prisma_client) if load_models else None + await self.get_credentials(prisma_client=prisma_client) + if load_models: still_desired_ids = await self._update_llm_router( new_models=new_models, proxy_logging_obj=proxy_logging_obj ) @@ -7148,12 +7160,9 @@ class ProxyConfig: async def _resync_config_from_db() -> None: await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) - async def _resync_credentials_from_db() -> None: - await self.get_credentials(prisma_client=prisma_client) - subscriber: Final = ConfigSyncSubscriber( redis_cache=redis_cache, - resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db), + resync_callbacks=(_resync_config_from_db,), ) self.config_sync_subscriber = subscriber subscriber.start() @@ -8008,7 +8017,7 @@ class ProxyConfig: async def get_credentials(self, prisma_client: PrismaClient): try: - credentials = await CredentialsRepository(prisma_client).find_all() + credentials = await CredentialsRepository(WriterPinnedClient(prisma_client.db)).find_all() credentials = [self.decrypt_credentials(cred) for cred in credentials] await self.delete_credentials(credentials) # delete credentials that are not in the all-up list CredentialAccessor.upsert_credentials(credentials) # upsert credentials that are in the all-up list @@ -9592,19 +9601,6 @@ class ProxyStartupEvent: ) if store_model_in_db is True: - ### GET STORED CREDENTIALS ### - scheduler.add_job( - proxy_config.get_credentials, - "interval", - seconds=config_reload_interval_seconds, - # REMOVED jitter parameter - major cause of memory leak - args=[prisma_client], - id="get_credentials_job", - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - await proxy_config.get_credentials(prisma_client=prisma_client) - # MEMORY LEAK FIX: Increase interval from 10s to 30s minimum # Frequent polling was causing excessive memory allocations scheduler.add_job( @@ -9618,7 +9614,7 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - # this will load all existing models on proxy startup + # this will load all existing credentials and models on proxy startup await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) proxy_config.start_config_sync_subscriber( 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/schema.prisma b/litellm/proxy/schema.prisma index 06ac177cca4..ccbab0fef10 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1509,6 +1509,8 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) 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..950fcca2039 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): @@ -323,6 +299,8 @@ def compute_autorouter_savings( selected_info: ModelInfo | None = None, baseline_info: ModelInfo | None = None, cost_breakdown: Mapping[str, object] | None = None, + baseline_deployment_id: str | None = None, + selected_deployment_id: str | None = None, ) -> float: """Net dollars the router saved, or cost, by serving this request on ``selected_model``. @@ -358,11 +336,12 @@ def compute_autorouter_savings( selected: Final = _resolve_model(selected_model, selected_provider) if baseline is None or selected is None: return 0.0 - # Same model is only the same cost when it is also the same deployment. Two - # deployments of one model can carry different negotiated rates, and routing from - # the dear one to the cheap one is a real saving that short-circuiting on the model - # name alone reports as zero. - if baseline == selected: + same_target: Final = ( + baseline_deployment_id == selected_deployment_id + if baseline_deployment_id and selected_deployment_id + else baseline == selected + ) + if same_target: return 0.0 basis: Final = _pricing_basis(cost_breakdown) effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline) @@ -541,6 +520,8 @@ def autorouter_savings_for_request( selected_info=_effective_model_info(router_instance, model_id, model or ""), baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""), cost_breakdown=cost_breakdown, + baseline_deployment_id=baseline_id, + selected_deployment_id=model_id, ) classifier_cost: Final = classifier_cost_from_decision(decision) return gross if classifier_cost is None else gross - classifier_cost @@ -586,6 +567,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 +577,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 +606,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 +623,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 ddf31cb1d8a..75d7972c621 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -177,6 +177,9 @@ from litellm.types.mcp import ( ) from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams +from litellm.utils import ( + _add_custom_logger_callback_to_specific_event, # pyright: ignore[reportPrivateUsage] # only string-to-logger helper +) if TYPE_CHECKING: from mcp.types import CallToolResult @@ -857,6 +860,14 @@ class ProxyLogging: litellm.logging_callback_manager.add_litellm_async_success_callback(callback) litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + # Runs after load_config applied every litellm_settings key: logger __init__s read e.g. s3_callback_params + success_callbacks: Final = tuple(cb for cb in litellm.success_callback if isinstance(cb, str)) + failure_callbacks: Final = tuple(cb for cb in litellm.failure_callback if isinstance(cb, str)) + for callback in success_callbacks: + _add_custom_logger_callback_to_specific_event(callback, "success") + for callback in failure_callbacks: + _add_custom_logger_callback_to_specific_event(callback, "failure") + async def update_request_status(self, litellm_call_id: str, status: Literal["success", "fail"]): # only use this if slack alerting is being used if self.alerting is None: 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/responses/main.py b/litellm/responses/main.py index 5e74b7324b4..fa14eb5f3c4 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -177,7 +177,7 @@ async def aresponses_api_with_mcp( ( mcp_tools_with_litellm_proxy, other_tools, - ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform) # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata) @@ -236,6 +236,7 @@ async def aresponses_api_with_mcp( "timeout": timeout, "custom_llm_provider": custom_llm_provider, **kwargs, + "_skip_mcp_handler": True, } # Handle MCP streaming if requested @@ -898,13 +899,14 @@ def _responses_try_dispatch_mcp_gateway( custom_llm_provider: str | None, kwargs: dict[str, object], _is_async: bool, + skip_mcp_handler: bool, ) -> Any | None: """Return a response when MCP gateway handles the call; otherwise None.""" from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): + if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(tools=tools): return None mcp_call_kwargs: Final = { "input": input, @@ -1074,6 +1076,7 @@ def responses( litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj") litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) _is_async: Final = kwargs.pop("aresponses", False) is True + skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) use_chat_completions_api = _pop_use_chat_completions_api_kw(kwargs) client_headers: Final = kwargs.get("headers") @@ -1168,6 +1171,7 @@ def responses( custom_llm_provider=custom_llm_provider, kwargs=kwargs, _is_async=_is_async, + skip_mcp_handler=skip_mcp_handler, ) if _mcp_dispatch is not None: return _mcp_dispatch diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index a75b3768636..ae18d5f6f1b 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -106,7 +106,7 @@ async def acompletion_with_mcp( ( mcp_tools_with_litellm_proxy, other_tools, - ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + ) = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools(tools) if not mcp_tools_with_litellm_proxy: # No MCP tools, proceed with regular completion @@ -114,6 +114,7 @@ async def acompletion_with_mcp( model=model, messages=messages, tools=tools, + _skip_mcp_handler=True, **kwargs, ) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 15434bedbb7..a5021e2f777 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,6 +1,6 @@ import re import traceback -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.utils import ( + iter_known_server_prefixes, logging_safe_mcp_headers, split_server_prefix_from_name, strip_known_server_prefix, @@ -23,6 +24,7 @@ from litellm.types.llms.openai import ( ResponsesAPIStreamingResponse, ) from litellm.types.llms.openai import ToolParam as ResponsesToolParam +from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.utils import ( CallTypes, ChatCompletionMessageCustomToolCall, @@ -45,6 +47,7 @@ else: # NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling ToolParam: TypeAlias = Mapping[str, object] +SplitTools: TypeAlias = tuple[list[ToolParam], list[Any]] class MCPToolResult(TypedDict): @@ -56,14 +59,74 @@ class MCPToolResult(TypedDict): LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" -# Matches any URL whose path ends with /mcp/ — covers both root-path -# (http://host:port/mcp/name) and sub-path (http://host/base/mcp/name) proxy deployments. -# A false-positive match (e.g. an external URL that happens to end with /mcp/) results -# in a "server not found" error from the internal gateway, not a silent failure or data leak, -# so this broad pattern is intentional and preferred over anchoring to localhost only. _PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") +def _mcp_server_url(tool: ToolParam) -> str | None: + if not isinstance(tool, dict) or tool.get("type") != "mcp": + return None + server_url: Final = tool.get("server_url") + return server_url if isinstance(server_url, str) else None + + +def _names_gateway_explicitly(tool: ToolParam) -> bool: + return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) + + +def _proxy_path_mcp_name(tool: ToolParam) -> str | None: + server_url: Final = _mcp_server_url(tool) + match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) + return None if match is None else match.group(1) + + +def _registered_mcp_servers() -> Collection[MCPServer]: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + return global_mcp_server_manager.get_registry().values() + + +def _registry_serves(name: str, servers: Collection[MCPServer]) -> bool: + requested: Final = name.lower() + return any( + requested in (known.lower() for known in (*iter_known_server_prefixes(server), server.name)) + or name in (server.access_groups or ()) + for server in servers + ) + + +async def _toolset_exists(name: str) -> bool: + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + return False + return await global_mcp_server_manager.get_toolset_by_name_cached(prisma_client, name) is not None + except Exception as e: + verbose_logger.debug("Could not resolve '%s' as toolset: %s", name, e) + return False + + +async def _gateway_served_names( + names: Collection[str], + servers: Callable[[], Collection[MCPServer]] = _registered_mcp_servers, + toolset_exists: Callable[[str], Awaitable[bool]] = _toolset_exists, +) -> frozenset[str]: + registered: Final = tuple(servers()) if names else () + return frozenset([name for name in names if _registry_serves(name, registered) or await toolset_exists(name)]) + + +async def _served_mcp_path_names( + tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] +) -> frozenset[str]: + names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) + return await served_names(names) if names else frozenset[str]() + + class LiteLLM_Proxy_MCP_Handler: """ Helper class with static methods for MCP integration with Responses API. @@ -87,57 +150,41 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: - """ - Returns True if any MCP tool should be handled via the litellm proxy MCP gateway. - This includes tools with server_url="litellm_proxy" as well as URLs ending in /mcp/. - """ - if tools: - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "mcp": - server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): - return True - if isinstance(server_url, str) and _PROXY_MCP_PATH_RE.match(server_url): - return True - return False + """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in + /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" + return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) @staticmethod - def _parse_mcp_tools( + def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: + items: Final = tuple(tools or ()) + gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] + other_tools: Final[list[Any]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + return gateway_tools, other_tools + + @staticmethod + async def _split_mcp_tools( tools: Iterable[Mapping[str, object]] | None, - ) -> tuple[list[ToolParam], list[Any]]: - """ - Parse tools and separate MCP tools with litellm_proxy from other tools. + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> SplitTools: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( + [ + {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} + if (name := _proxy_path_mcp_name(tool)) in served + else tool + for tool in items + ] + ) - Returns: - Tuple of (mcp_tools_with_litellm_proxy, other_tools) - """ - mcp_tools_with_litellm_proxy: Final[list[ToolParam]] = [] - other_tools: Final[list[Any]] = [] - - if tools: - for tool in tools: - if isinstance(tool, dict) and tool.get("type") == "mcp": - server_url = tool.get("server_url", "") - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL): - mcp_tools_with_litellm_proxy.append(tool) - elif isinstance(server_url, str): - # Also intercept URLs like http://localhost:4000/mcp/atlassian_test - # by rewriting them to the internal litellm_proxy format. - m = _PROXY_MCP_PATH_RE.match(server_url) - if m: - rewritten = { - **tool, - "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{m.group(1)}", - } - mcp_tools_with_litellm_proxy.append(rewritten) - else: - other_tools.append(tool) - else: - other_tools.append(tool) - else: - other_tools.append(tool) - - return mcp_tools_with_litellm_proxy, other_tools + @staticmethod + async def routes_through_gateway( + tools: Iterable[Mapping[str, object]] | None, + served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] = _gateway_served_names, + ) -> tuple[bool, ...]: + items: Final = tuple(tools or ()) + served: Final = await _served_mcp_path_names(items, served_names) + return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) @staticmethod async def _apply_toolset_permissions( diff --git a/litellm/router.py b/litellm/router.py index c7a254cbc02..9b2badeaa69 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -21,7 +21,16 @@ import time import traceback import weakref from collections import defaultdict -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Callable, + Generator, + Iterator, + Mapping, + MutableMapping, + Sequence, +) from functools import lru_cache, partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast @@ -45,12 +54,14 @@ from litellm.caching.caching import ( RedisClusterCache, ) from litellm.constants import ( + CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, DEFAULT_HEALTH_CHECK_INTERVAL, DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, + OUTPUT_TOKEN_CEILING_PARAMS, RUNTIME_UPDATABLE_ROUTER_SETTINGS, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) @@ -616,6 +627,50 @@ 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 + + +def as_output_cap(value: object) -> int | None: + """A client-sent output cap coerced to an int: ints, floats and numeric strings, never bools + or negatives.""" + if isinstance(value, bool) or not isinstance(value, (int, float, str)): + return None + try: + cap: Final = int(float(value)) + except (ValueError, OverflowError): + return None + return cap if cap >= 0 else None + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -923,7 +978,6 @@ class Router: DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER ) self.health_state_cache = DeploymentHealthCache(cache=self.cache, staleness_threshold=float(_staleness)) - self.failed_calls = InMemoryCache() # cache to track failed call per deployment, if num failed calls within 1 minute > allowed fails, then add it to cooldown if num_retries is not None: self.num_retries = num_retries @@ -1216,7 +1270,7 @@ class Router: selector = LeastBusyLoggingHandler(router_cache=self.cache) if register_callbacks: if isinstance(litellm.input_callback, list): - litellm.input_callback.append(selector) + litellm.logging_callback_manager.add_litellm_input_callback(selector) else: litellm.input_callback = [selector] case RoutingStrategy.USAGE_BASED_ROUTING.value: @@ -1491,9 +1545,33 @@ class Router: self._override_selectors[strategy] = self._build_strategy_selector( strategy=strategy, routing_strategy_args={}, + register_callbacks=False, ) return self._override_selectors[strategy] + def _override_selector_pre_call_check( + self, strategy: str | None, selector: RouterStrategySelector | None, deployment: dict + ) -> None: + """ + Override selectors are not in `litellm.callbacks`, so the pre-call check that + `routing_strategy_pre_call_checks` runs for the router's own selectors (rpm + accounting for `usage-based-routing-v2`) runs here, for the overriding request only. + """ + if selector is None or strategy is None or selector is not self._override_selectors.get(strategy): + return + selector.pre_call_check(deployment) + + async def _async_override_selector_pre_call_check( + self, + strategy: str | None, + selector: RouterStrategySelector | None, + deployment: dict, + parent_otel_span: Span | None, + ) -> None: + if selector is None or strategy is None or selector is not self._override_selectors.get(strategy): + return + await selector.async_pre_call_check(deployment, parent_otel_span) + def _get_routing_context( self, model: str, request_kwargs: dict | None = None ) -> tuple[str | None, RouterStrategySelector | None]: @@ -2576,6 +2654,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, @@ -2615,7 +2705,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__( @@ -2623,6 +2713,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) @@ -2699,8 +2790,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 @@ -2742,7 +2842,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( @@ -3171,13 +3275,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"): @@ -3233,8 +3338,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 @@ -3272,7 +3386,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): """ @@ -4143,10 +4260,12 @@ class Router: } ) litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) - prompt_management_deployment: Final = self.get_available_deployment( + specific_deployment: Final = kwargs.pop("specific_deployment", None) + prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), + messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally + specific_deployment=specific_deployment, + request_kwargs=kwargs, ) self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs) @@ -12533,7 +12652,83 @@ class Router: request_kwargs.pop(carrier, None) @staticmethod - def _drop_client_effort_carriers_a_tier_pin_supersedes( + def _tier_ceiling_under_the_surface_name( + tier_litellm_params: Mapping[str, object], responses_call: bool + ) -> Mapping[str, object]: + """``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are one + ceiling under three names, and each surface reads exactly one of them: the + Responses bridge builds its internal ``max_tokens`` from ``max_output_tokens`` + and would overwrite the tier's, chat and /v1/messages never read + ``max_output_tokens``, and litellm already renames ``max_tokens`` to + ``max_completion_tokens`` for the OpenAI models that require it. Collapse + whatever the tier carries onto the surface's own name, preferring a value the + operator already wrote under that name.""" + surface_key: Final = "max_output_tokens" if responses_call else "max_tokens" + carried: Final = tuple( + key + for key in (surface_key, "max_tokens", "max_completion_tokens", "max_output_tokens") + if key in tier_litellm_params + ) + if not carried: + return tier_litellm_params + return MappingProxyType( + { + **{k: v for k, v in tier_litellm_params.items() if k not in OUTPUT_TOKEN_CEILING_PARAMS}, + surface_key: tier_litellm_params[carried[0]], + } + ) + + def _pin_tier_params_onto_request( + self, + model: str, + tier_litellm_params: Mapping[str, object] | None, + request_kwargs: dict, + responses_call: bool, + ) -> bool: + """Apply a routing strategy's per-tier litellm_params on top of the request and report + whether they pinned an output ceiling, so the caller can hand the request its own ceiling + back on a routing pass that pins none.""" + if not tier_litellm_params: + return False + accepted_tier_params: Final = self._tier_params_the_target_accepts(model, tier_litellm_params, request_kwargs) + surface_tier_params: Final = self._tier_ceiling_under_the_surface_name( + accepted_tier_params, responses_call=responses_call + ) + self._drop_client_carriers_a_tier_pin_supersedes(request_kwargs, surface_tier_params) + request_kwargs.update(surface_tier_params) + return not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(surface_tier_params) + + @staticmethod + def _restore_client_ceiling_no_tier_pins(request_kwargs: MutableMapping[str, object]) -> None: + """A model-group fallback re-enters routing with the kwargs an earlier auto-router pass + already rewrote, so a ceiling sized for that pass's tier would ride onto a group no tier + chose. When this pass pins none, hand the request back exactly the carriers the caller + sent, which the first pinning pass stamped. The stamp lives in a metadata bucket a + caller can also write, so the proxy strips the key at ingestion and this read takes + nothing but the three ceiling carriers as integers: no other key ever reaches kwargs.""" + stamped: Final = next( + ( + bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY) + for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) + if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket + ), + None, + ) + if not isinstance(stamped, dict): + return + callers_ceiling: Final = MappingProxyType( + { + carrier: cap + for carrier, value in stamped.items() + if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None + } + ) + for carrier in OUTPUT_TOKEN_CEILING_PARAMS: + request_kwargs.pop(carrier, None) + request_kwargs.update(callers_ceiling) + + @staticmethod + def _drop_client_carriers_a_tier_pin_supersedes( request_kwargs: dict[str, object], tier_litellm_params: Mapping[str, object], ) -> None: @@ -12543,7 +12738,22 @@ class Router: the ``reasoning_effort`` alias, so a pinned effort only reaches the wire if the client's other encodings are removed before the merge. Non-effort fields a carrier also holds (``output_config.format``, - ``reasoning.summary``) are kept.""" + ``reasoning.summary``) are kept. An output ceiling has the same shape: + ``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are + one setting under three names, and a provider handed two of them either + rejects the request or picks one by iteration order.""" + if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): + _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) + metadata_bucket.setdefault( + CLIENT_OUTPUT_CEILING_METADATA_KEY, + { + carrier: request_kwargs[carrier] + for carrier in OUTPUT_TOKEN_CEILING_PARAMS + if carrier in request_kwargs + }, + ) + for carrier in OUTPUT_TOKEN_CEILING_PARAMS: + request_kwargs.pop(carrier, None) if "reasoning_effort" not in tier_litellm_params: return request_kwargs.pop("thinking", None) @@ -12584,6 +12794,7 @@ class Router: # Execute Pre-Routing Hooks # this hook can modify the model, messages before the routing decision is made ######################################################### + responses_call: Final = input is not None and messages is None pre_routing_hook_response: Final = await self.async_pre_routing_hook( model=model, request_kwargs=request_kwargs, @@ -12595,12 +12806,14 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - if pre_routing_hook_response.litellm_params: - accepted_tier_params: Final = self._tier_params_the_target_accepts( - model, pre_routing_hook_response.litellm_params, request_kwargs - ) - self._drop_client_effort_carriers_a_tier_pin_supersedes(request_kwargs, accepted_tier_params) - request_kwargs.update(accepted_tier_params) + tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + model=model, + tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, + request_kwargs=request_kwargs, + responses_call=responses_call, + ) + if not tier_pins_ceiling: + self._restore_client_ceiling_no_tier_pins(request_kwargs) ######################################################### # Resolve the strategy and logger AFTER the pre-routing hook, since @@ -12617,10 +12830,16 @@ class Router: parent_otel_span=parent_otel_span, ) if isinstance(healthy_deployments, dict): + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, healthy_deployments, parent_otel_span + ) return healthy_deployments # When encrypted content affinity pins to a specific deployment, if request_kwargs.get("_encrypted_content_affinity_pinned") and len(healthy_deployments) == 1: + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, healthy_deployments[0], parent_otel_span + ) return healthy_deployments[0] start_time: Final = time.time() @@ -12646,6 +12865,9 @@ class Router: parent_otel_span=parent_otel_span, ) raise exception + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, deployment, parent_otel_span + ) verbose_router_logger.info( "get_available_deployment for model: %s, Selected deployment: %s for model: %s", model, @@ -12700,6 +12922,7 @@ class Router: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs) # 1. Execute pre-routing hook + responses_call: Final = input is not None and messages is None pre_routing_hook_response: Final = await self.async_pre_routing_hook( model=model, request_kwargs=request_kwargs, @@ -12711,12 +12934,14 @@ class Router: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages record_pre_routing_selection(request_kwargs, model) - if pre_routing_hook_response.litellm_params: - accepted_tier_params: Final = self._tier_params_the_target_accepts( - model, pre_routing_hook_response.litellm_params, request_kwargs - ) - self._drop_client_effort_carriers_a_tier_pin_supersedes(request_kwargs, accepted_tier_params) - request_kwargs.update(accepted_tier_params) + tier_pins_ceiling: Final = self._pin_tier_params_onto_request( + model=model, + tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, + request_kwargs=request_kwargs, + responses_call=responses_call, + ) + if not tier_pins_ceiling: + self._restore_client_ceiling_no_tier_pins(request_kwargs) # 2. Get healthy deployments healthy_deployments: Final = await self.async_get_healthy_deployments( @@ -12728,6 +12953,8 @@ class Router: parent_otel_span=parent_otel_span, ) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) + # 3. If specific deployment returned, verify if it supports pass-through if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -12738,6 +12965,9 @@ class Router: ) litellm_params: Final = healthy_deployments.get("litellm_params", {}) if litellm_params.get("use_in_pass_through"): + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, healthy_deployments, parent_otel_span + ) return healthy_deployments else: raise litellm.BadRequestError( @@ -12758,7 +12988,6 @@ class Router: # 5. Apply load balancing strategy start_time: Final = time.perf_counter() - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": return simple_shuffle( llm_router_instance=self, @@ -12782,6 +13011,9 @@ class Router: parent_otel_span=parent_otel_span, ) raise exception + await self._async_override_selector_pre_call_check( + strategy, strategy_selector, deployment, parent_otel_span + ) verbose_router_logger.info( "async_get_available_deployment_for_pass_through model: %s, selected deployment: %s", @@ -13357,6 +13589,7 @@ class Router: specific_deployment=specific_deployment, request_kwargs=request_kwargs, ) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -13365,6 +13598,7 @@ class Router: model=model, llm_provider="", ) + self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments) return healthy_deployments parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(request_kwargs) @@ -13440,7 +13674,6 @@ class Router: cooldown_list=_cooldown_list, ) - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": # if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm ############## Check 'weight' param set for weighted pick ################# @@ -13472,6 +13705,7 @@ class Router: enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, ) + self._override_selector_pre_call_check(strategy, strategy_selector, deployment) verbose_router_logger.info( "get_available_deployment for model: %s, Selected deployment: %s for model: %s", model, @@ -13515,6 +13749,8 @@ class Router: specific_deployment=specific_deployment, ) + strategy, strategy_selector = self._get_routing_context(model, request_kwargs) + # 2. If the returned is a specific deployment (Dict), verify and return directly if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -13525,6 +13761,7 @@ class Router: ) litellm_params: Final = healthy_deployments.get("litellm_params", {}) if litellm_params.get("use_in_pass_through"): + self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments) return healthy_deployments else: # Specific deployment does not support pass-through @@ -13584,7 +13821,6 @@ class Router: ) # 6. Apply load balancing strategy - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) if strategy == "simple-shuffle": return simple_shuffle( llm_router_instance=self, @@ -13616,6 +13852,7 @@ class Router: enable_pre_call_checks=self.enable_pre_call_checks, cooldown_list=_cooldown_list, ) + self._override_selector_pre_call_check(strategy, strategy_selector, deployment) verbose_router_logger.info( "get_available_deployment_for_pass_through model: %s, selected deployment: %s", diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 93dddfb3d20..d605b43e42a 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -195,6 +195,40 @@ model_list: session_affinity_ttl_seconds: 300 ``` +## Custom dimensions + +Add `custom_dimensions` under `complexity_router_config` to give domain keywords or regex patterns their own weighted signal + +```yaml +custom_dimensions: + - name: internalFrameworks + weight: 0.9 + keywords: [orbitmesh, fluxgate] + - name: sqlMigration + weight: 0.7 + patterns: ['\b(create|alter|drop)\s{1,4}table\b'] + - name: dataPipeline + weight: 0.4 + scoring_mode: match_count + keywords: [airflow, dbt, snowflake] +``` + +Each dimension contributes its weight once when any matcher hits the current ask. Repeated matches do not increase it. The built-in score and tier boundaries are unchanged, and the total score is not renormalized. Keywords use the existing case-insensitive word-boundary and CJK rules. Regexes search the first 2048 characters case-insensitively and compile during configuration validation and router initialization, never per request + +`scoring_mode` is optional and defaults to `binary`, the behavior above. `match_count` grades the dimension by how many distinct matchers hit: none scores 0 and emits no signal, one scores half the weight, two or more score the full weight. Repeated occurrences of one matcher never raise the count, keywords are distinct case-insensitively, patterns are distinct by source, and a keyword and a pattern are always distinct from each other. Matching stops as soon as the selected mode's maximum is reached, so a binary dimension still stops at its first hit. Existing configurations without the field keep binary scoring and the same tuning fingerprint, so the field only counts as a tuning change when set to `match_count` + +### Weights through the API versus the dashboard + +The API and YAML store exactly the weights written. A `dimension_weights` map and inline custom weights are read literally, missing recognized built-in names score zero, and nothing renormalizes the vector, so a total other than 1 is legal and scores accordingly. The dashboard's heuristic scoring editor is the one place that rebalances: editing one weight there holds it and redistributes the remainder across the other active dimensions in the draft, then Save sends the resulting explicit values, which the backend stores and scores as written. Opening a router, applying a preset, editing matchers, changing `scoring_mode`, or saving unrelated fields never normalizes existing weights + +Only `heuristic`, `heuristic_first` and `hybrid` accept custom dimensions. Each name must be a unique ASCII identifier starting with a letter, at most 64 characters, and cannot reuse a built-in dimension name or a key in `dimension_weights`. Set its weight inline, greater than zero and at most one + +Patterns are checked at configuration time against a grammar whose worst case stays a few milliseconds on 2048 characters. Every quantifier needs an explicit upper bound of at most 64 and must repeat a single character or character class, so `\s{1,4}` is accepted while `\s+`, `(a|aa){0,12}` and `(?:ab){0,64}` are refused. Backreferences, lookarounds, atomic groups and possessive quantifiers are refused as well. Each pattern is then costed: alternation branches and repeat lengths multiply the ways the engine can retry, and every later piece of the pattern is charged once per path that can reach it, so `a?a?a?a?a?a?a?a?` followed by a long fixed tail is refused even though each quantifier is small. The budget is 2048 work units per pattern and 8192 across the router. An invalid or over-budget pattern fails the write with a message naming the pattern and the rule it broke + +Limits are 16 dimensions, 32 combined keywords/patterns per dimension, 256 characters per matcher and 4096 matcher characters per dimension. Matching runs inline on the request path with no timeout and no worker thread, because the grammar is what bounds the cost. These are routing hints, not security enforcement rules + +The existing heuristic-v1 tuning quota covers custom dimensions, their weights and their scoring mode: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML, the model API, or the dashboard's heuristic scoring editor + ## Usage Once configured, use the model name like any other: diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 7d4497fb6f7..f9b4f50c15c 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -20,7 +20,7 @@ import random import re import time from collections.abc import Callable, Iterator, Mapping, Sequence -from itertools import accumulate, islice, takewhile +from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast @@ -31,6 +31,7 @@ from litellm._logging import verbose_router_logger from litellm.constants import ( EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, + OUTPUT_TOKEN_CEILING_PARAMS, RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) @@ -67,6 +68,7 @@ from litellm.types.utils import ( from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section from .config import ( CALIBRATION_EXAMPLES_HEADING, + CUSTOM_PATTERN_SCAN_CHARS, DEFAULT_CLASSIFICATION_RUBRIC, DEFAULT_CODE_KEYWORDS, DEFAULT_ESCALATION_KEYWORDS, @@ -81,6 +83,7 @@ from .config import ( ClassificationRubric, ComplexityRouterConfig, ComplexityTier, + CustomDimension, TierDefinition, ) from .stall_detector import detect_stalled_task @@ -878,6 +881,15 @@ class DimensionScore: self.signal = signal +class _CustomDimensionMatchers(NamedTuple): + """One custom dimension's distinct matchers and the number of hits that saturates its score.""" + + dimension: CustomDimension + keywords: tuple[str, ...] + patterns: tuple[re.Pattern[str], ...] + saturation: int + + class KeywordOverride(NamedTuple): """A keyword_tier_rules match: the winning tier and, on the lexical path, the keyword that fired.""" @@ -1119,6 +1131,15 @@ class ComplexityRouter(CustomLogger): self.config.custom_technical_keywords, ) self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS + self._custom_dimensions = tuple( + _CustomDimensionMatchers( + dimension, + tuple(dict.fromkeys(keyword.lower() for keyword in dimension.keywords)), + tuple(re.compile(pattern, re.IGNORECASE) for pattern in dict.fromkeys(dimension.patterns)), + 2 if dimension.scoring_mode == "match_count" else 1, + ) + for dimension in self.config.custom_dimensions + ) if self.config.has_custom_tiers: self.escalation_keywords: tuple[str, ...] = () elif self.config.escalation_keywords is not None: @@ -1320,6 +1341,28 @@ class ComplexityRouter(CustomLogger): score: Final = score_high if match_count >= high_threshold else score_low return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count + def _count_custom_hits(self, matchers: _CustomDimensionMatchers, user_text: str, scanned: str) -> int: + hits: Final = chain( + (self._keyword_matches(user_text, keyword) for keyword in matchers.keywords), + (pattern.search(scanned) is not None for pattern in matchers.patterns), + ) + return sum(islice((1 for hit in hits if hit), matchers.saturation)) + + def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]: + if not self._custom_dimensions: + return () + scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS] + return tuple( + ( + DimensionScore( + matchers.dimension.name, hits / matchers.saturation, f"custom ({matchers.dimension.name})" + ), + matchers.dimension.weight, + ) + for matchers in self._custom_dimensions + if (hits := self._count_custom_hits(matchers, user_text, scanned)) + ) + def _score_multi_step(self, text: str) -> DimensionScore: """Score based on multi-step patterns.""" hits: Final = sum(1 for p in self._multi_step_patterns if p.search(text)) @@ -1415,12 +1458,13 @@ class ComplexityRouter(CustomLogger): self._score_question_complexity(prompt), ] - # Collect signals - signals: Final = [d.signal for d in dimensions if d.signal is not None] + custom_dimensions: Final = self._score_custom_dimensions(prompt, user_text) + signals: Final = [d.signal for d in (*dimensions, *(d for d, _ in custom_dimensions)) if d.signal is not None] - # Compute weighted score weights: Final = self.config.dimension_weights - weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + weighted_score: Final = sum(d.score * weights.get(d.name, 0) for d in dimensions) + sum( + dimension.score * weight for dimension, weight in custom_dimensions + ) boundaries: Final = self._effective_tier_boundaries() clears_override_floor: Final = weighted_score >= self._effective_reasoning_override_min_score() @@ -2067,11 +2111,15 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"No model configured for tier {tier_key} and no default_model set") def _litellm_params_for_model(self, tier: ComplexityTier | str | None, model: str) -> Mapping[str, object]: - if tier is None: - return MappingProxyType({}) - entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) + entries: Final = self.config.tier_model_configs.get(_tier_name(tier), ()) if tier is not None else () entry: Final = next((candidate for candidate in entries if candidate.model_name == model), None) - return entry.litellm_params if entry is not None else MappingProxyType({}) + explicit: Final = entry.litellm_params if entry is not None else MappingProxyType({}) + if not self.config.max_tokens_from_tier_model or not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(explicit): + return explicit + ceiling: Final = self._group_output_ceiling(model) + if ceiling is None: + return explicit + return MappingProxyType({**explicit, "max_tokens": ceiling}) @staticmethod def _pick_from_tier_value(model: str | Sequence[str], tier_key: str) -> str: @@ -2406,12 +2454,15 @@ class ComplexityRouter(CustomLogger): return name if self.config.has_custom_tiers else ComplexityTier(name) def _deployment_window(self, group: str, deployment: Mapping[str, object]) -> int | None: + return self._deployment_limit(group, deployment, "max_input_tokens") + + def _deployment_limit( + self, group: str, deployment: Mapping[str, object], key: Literal["max_input_tokens", "max_output_tokens"] + ) -> int | None: from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider deployment_model_info: Final = deployment.get("model_info") - declared: Final = ( - deployment_model_info.get("max_input_tokens") if isinstance(deployment_model_info, Mapping) else None - ) + declared: Final = deployment_model_info.get(key) if isinstance(deployment_model_info, Mapping) else None if isinstance(declared, int): return declared litellm_params: Final = deployment.get("litellm_params") @@ -2428,18 +2479,34 @@ class ComplexityRouter(CustomLogger): deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts received_model_name=group, ) - window: Final = model_info.get("max_input_tokens") + limit: Final = model_info.get(key) except Exception: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others return None - return window if isinstance(window, int) else None + return limit if isinstance(limit, int) else None + + def _group_deployments(self, group: str) -> Sequence[Mapping[str, object]]: + list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None) + deployments: Final = list_models(model_name=group) if callable(list_models) else None + return tuple(deployments) if isinstance(deployments, list) else () + + def _group_output_ceiling(self, group: str) -> int | None: + """Smallest max_output_tokens across the group's deployments, or None when any deployment + declares none: the core router picks within the group without a fit check, and a ceiling + above an unmapped member's real limit is a provider 400 on that member.""" + deployments: Final = self._group_deployments(group) + ceilings: Final = tuple( + ceiling + for deployment in deployments + if (ceiling := self._deployment_limit(group, deployment, "max_output_tokens")) is not None + ) + return min(ceilings) if ceilings and len(ceilings) == len(deployments) else None def _group_window_facts(self, group: str) -> tuple[int | None, bool]: """(smallest declared context window across the group's deployments, whether any deployment declares none). The core router picks a deployment within the group without a fit check, so the group is only as safe as its smallest member.""" - list_models: Final = getattr(self.litellm_router_instance, "get_model_list", None) - deployments: Final = list_models(model_name=group) if callable(list_models) else None - if not isinstance(deployments, list) or not deployments: + deployments: Final = self._group_deployments(group) + if not deployments: return (None, True) windows: Final = tuple( window for deployment in deployments if (window := self._deployment_window(group, deployment)) is not None @@ -3488,6 +3555,7 @@ class ComplexityRouter(CustomLogger): ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs ) fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM + default_tier_params: Final = self._litellm_params_for_model(fallback_tier, routed_model) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -3496,7 +3564,9 @@ class ComplexityRouter(CustomLogger): cause="default_fallback", tier=fallback_tier, conversation_continuing=conversation_continuing, + tier_litellm_params=default_tier_params, ), + litellm_params=default_tier_params, ) ask: Final = user_message or "" @@ -3523,6 +3593,7 @@ class ComplexityRouter(CustomLogger): _tier_name(plan_floor), routed_model, ) + plan_tier_params: Final = self._litellm_params_for_model(plan_floor, routed_model) return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, @@ -3534,7 +3605,9 @@ class ComplexityRouter(CustomLogger): matched_keyword=plan_mode_sentinel, escalation_keyword=escalation_keyword, escalated=False, + tier_litellm_params=plan_tier_params, ), + litellm_params=plan_tier_params, ) override: Final = await self._resolve_keyword_tier_override(ask, request_kwargs) @@ -3627,6 +3700,7 @@ class ComplexityRouter(CustomLogger): outcome.signals, fallback_model, ) + fallback_tier_params: Final = self._litellm_params_for_model(None, fallback_model) return PreRoutingHookResponse( model=fallback_model, messages=messages if has_original_messages else None, @@ -3637,7 +3711,9 @@ class ComplexityRouter(CustomLogger): signals=outcome.signals, escalation_keyword=escalation_keyword, escalated=False, + tier_litellm_params=fallback_tier_params, ), + litellm_params=fallback_tier_params, ) if self.config.adaptive: # hard_floor rather than a hard pick, and passed whenever the sentinel is present diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index c483a0b7073..58f56d746c5 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -5,13 +5,21 @@ Contains default keyword lists, weights, tier boundaries, and configuration clas All values are configurable via proxy config.yaml. """ -from collections.abc import Mapping +import math +import re +import warnings +from collections.abc import Iterable, Mapping from enum import Enum from types import MappingProxyType -from typing import Annotated, Final, Literal +from typing import Annotated, Final, Literal, NamedTuple from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator +with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + import sre_constants + import sre_parse + from litellm.types.llms.openai import REASONING_EFFORT from litellm.types.router import AdaptiveRouterWeights, ClassifierPlugin, RoutingPlugin @@ -569,6 +577,125 @@ class ClassifierLLMConfig(BaseModel): return self +MAX_CUSTOM_PATTERN_REPEAT: Final[int] = 64 +MAX_CUSTOM_PATTERN_WORK: Final[int] = 2048 +MAX_CUSTOM_DIMENSIONS_WORK: Final[int] = 8192 +MAX_CUSTOM_PATTERN_DEPTH: Final[int] = 16 +CUSTOM_PATTERN_SCAN_CHARS: Final[int] = 2048 + +_ATOM_OPCODES: Final = frozenset( + {sre_constants.LITERAL, sre_constants.NOT_LITERAL, sre_constants.ANY, sre_constants.IN, sre_constants.CATEGORY} +) +_REPEAT_OPCODES: Final = frozenset({sre_constants.MAX_REPEAT, sre_constants.MIN_REPEAT}) + + +class _PatternCost(NamedTuple): + paths: int + steps: int + + +def _atom_steps(node: object) -> int: + if isinstance(node, tuple) and len(node) == 2 and node[0] is sre_constants.IN: + return 1 + len(node[1]) + return 1 + + +def _repeat_cost(argument: object) -> _PatternCost | str: + if not isinstance(argument, tuple) or len(argument) != 3: + return "unsupported repeat structure" + low, high, body = argument + if high > MAX_CUSTOM_PATTERN_REPEAT or len(body) != 1 or body[0][0] not in _ATOM_OPCODES: + return "requires a single character or class repeated at most 64 times; use {n,m} instead of *, + or {n,}" + choices: Final = high - low + 1 + return _PatternCost(choices, 1 + high * _atom_steps(body[0]) + choices) + + +def _node_cost(node: object, depth: int) -> _PatternCost | str: + if not isinstance(node, tuple) or len(node) != 2: + return "unsupported regex structure" + opcode, argument = node + if opcode in _ATOM_OPCODES or opcode is sre_constants.AT: + return _PatternCost(1, _atom_steps(node)) + if opcode is sre_constants.SUBPATTERN: + return _sequence_cost(argument[-1], depth + 1) + if opcode is sre_constants.BRANCH: + costs: Final = tuple(_sequence_cost(branch, depth + 1) for branch in argument[1]) + refused: Final = next((cost for cost in costs if isinstance(cost, str)), None) + if refused is not None: + return refused + return _PatternCost( + sum(cost.paths for cost in costs if isinstance(cost, _PatternCost)), + len(costs) + sum(cost.steps for cost in costs if isinstance(cost, _PatternCost)), + ) + if opcode in _REPEAT_OPCODES: + return _repeat_cost(argument) + return "contains an unsupported regex construct" + + +def _sequence_cost(nodes: Iterable[object], depth: int) -> _PatternCost | str: + if depth > MAX_CUSTOM_PATTERN_DEPTH: + return "nests deeper than 16 levels" + costs: Final = tuple(_node_cost(node, depth) for node in nodes) + refused: Final = next((cost for cost in costs if isinstance(cost, str)), None) + if refused is not None: + return refused + valid: Final = tuple(cost for cost in costs if isinstance(cost, _PatternCost)) + # Choices multiply across a sequence; every continuation can execute once per preceding path. + total: Final = _PatternCost( + math.prod(cost.paths for cost in valid), + 1 + sum(cost.steps * math.prod(prior.paths for prior in valid[:index]) for index, cost in enumerate(valid)), + ) + if total.steps > MAX_CUSTOM_PATTERN_WORK: + return "exceeds the per-pattern regex work budget" + return total + + +def custom_pattern_work(pattern: str) -> int | str: + try: + re.compile(pattern, re.IGNORECASE) + parsed: Final = sre_parse.parse(pattern, re.IGNORECASE) + except (re.error, RecursionError, OverflowError): + return "is not a valid regex" + cost: Final = _sequence_cost(tuple(parsed), 0) + return cost if isinstance(cost, str) else cost.steps + + +class CustomDimension(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + name: str = Field(min_length=1, max_length=64, pattern=r"^[A-Za-z][A-Za-z0-9_]*$") + weight: float = Field(gt=0, le=1, allow_inf_nan=False) + keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32) + patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32) + scoring_mode: Literal["binary", "match_count"] = Field( + default="binary", + description=( + "'binary' scores 1 when any matcher hits. 'match_count' scores 0.5 when one distinct matcher hits and 1 " + "when two or more do; repeated occurrences of one matcher never raise it. Keywords are distinct " + "case-insensitively, patterns by source, and a keyword and a pattern are always distinct from each other." + ), + ) + + @model_validator(mode="after") + def _validate_matchers(self) -> "CustomDimension": + matchers: Final = (*self.keywords, *self.patterns) + if not matchers or any(not matcher.strip() for matcher in matchers): + raise ValueError("custom dimensions require nonblank keywords and/or patterns") + if len(matchers) > 32 or sum(map(len, matchers)) > 4096: + raise ValueError("custom dimensions allow at most 32 matchers and 4096 matcher characters each") + costs: Final = tuple((pattern, custom_pattern_work(pattern)) for pattern in self.patterns) + rejected: Final = tuple(f"pattern {pattern!r} {work}" for pattern, work in costs if isinstance(work, str)) + if rejected: + raise ValueError("custom dimension " + "; ".join(rejected)) + return self + + def pattern_work(self) -> int: + """Combined work estimate of the validated patterns.""" + return sum( + work for work in (custom_pattern_work(pattern) for pattern in self.patterns) if isinstance(work, int) + ) + + class ComplexityRouterConfig(BaseModel): """Configuration for the ComplexityRouter.""" @@ -671,6 +798,20 @@ class ComplexityRouterConfig(BaseModel): description="Weights for each scoring dimension", ) + custom_dimensions: tuple[CustomDimension, ...] = Field( + default=(), + max_length=16, + description=( + "Named dimensions added to the heuristic-v1 score. Each contributes its inline weight once " + "when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters; " + "scoring_mode 'match_count' instead grades half weight for one distinct matcher and full for two or more. " + "Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, " + "backreferences and lookarounds are rejected. Conservative work limits include alternation paths, " + "repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. " + "Only heuristic, heuristic_first and hybrid accept this field. Uses the existing heuristic tuning quota." + ), + ) + # Keyword lists (overridable) code_keywords: list[str] | None = Field( default=None, @@ -948,6 +1089,20 @@ class ComplexityRouterConfig(BaseModel): "wording the built-ins don't cover, or after a client release changes its strings." ), ) + max_tokens_from_tier_model: bool = Field( + default=True, + description=( + "Set max_tokens on every routed request to the output ceiling of the tier model it " + "lands on, replacing whatever the caller sent. A caller behind an auto-router cannot " + "pick one value that fits every tier: the smallest tier's ceiling starves a bigger " + "tier's thinking budget, and a bigger tier's ceiling is rejected by the smallest. The " + "ceiling is the smallest max_output_tokens across the tier model's deployments, read " + "from each deployment's model_info and then the model cost map; a tier model with a " + "deployment whose ceiling is unknown keeps the caller's value. A max_tokens, " + "max_completion_tokens or max_output_tokens in the tier's own litellm_params still " + "wins. Set false to forward the caller's value unchanged." + ), + ) route_housekeeping_to_cheapest_tier: bool = Field( default=True, description=( @@ -1245,6 +1400,27 @@ class ComplexityRouterConfig(BaseModel): ) return self + @model_validator(mode="after") + def _validate_custom_dimensions(self) -> "ComplexityRouterConfig": + if not self.custom_dimensions: + return self + if self.classifier_type not in ("heuristic", "heuristic_first", "hybrid"): + raise ValueError("custom_dimensions requires classifier_type heuristic, heuristic_first or hybrid") + names: Final = tuple(dimension.name.casefold() for dimension in self.custom_dimensions) + reserved: Final = frozenset(name.casefold() for name in DEFAULT_DIMENSION_WEIGHTS) + weighted: Final = frozenset(name.casefold() for name in self.dimension_weights) + if len(frozenset(names)) != len(names) or frozenset(names) & reserved: + raise ValueError("custom dimension names must be unique and must not shadow built-in dimensions") + if frozenset(names) & weighted: + raise ValueError("custom dimension weights must be inline, not in dimension_weights") + work: Final = sum(dimension.pattern_work() for dimension in self.custom_dimensions) + if work > MAX_CUSTOM_DIMENSIONS_WORK: + raise ValueError( + f"custom_dimensions regex work estimate is {work}; the limit across the router is " + f"{MAX_CUSTOM_DIMENSIONS_WORK}" + ) + return self + @field_validator("heuristic_first_max_tier", mode="before") @classmethod def _coerce_heuristic_first_max_tier(cls, value: object) -> object: diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1433e8ba4d4..14e6592e1fd 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,17 +1,103 @@ -#### What this does #### -# identifies least busy deployment -# How is this achieved? -# - Before each call, have the router print the state of requests {"deployment": "requests_in_flight"} -# - use litellm.input_callbacks to log when a request is just about to be made to a model - {"deployment-id": traffic} -# - use litellm.success + failure callbacks to log when a request completed -# - in get_available_deployment, for a given model group name -> pick based on traffic - -import random +from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 + + +class _ModelInfo(TypedDict, total=False): + id: ReadOnly[str | int | None] + + +class _Metadata(TypedDict, total=False): + model_group: ReadOnly[str | None] + + +class _LitellmParams(TypedDict, total=False): + metadata: ReadOnly[_Metadata | None] + model_info: ReadOnly[_ModelInfo | None] + + +class _CallKwargs(TypedDict, total=False): + litellm_params: ReadOnly[_LitellmParams | None] + + +class _DeploymentModelInfo(TypedDict): + id: ReadOnly[str | int] + + +class _Deployment(TypedDict): + model_info: ReadOnly[_DeploymentModelInfo] + + +_CALL_KWARGS: Final = TypeAdapter(_CallKwargs) +_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) +_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) + + +def _request_count_key(model_group: str, deployment_id: str) -> str: + return f"{model_group}_request_count:{deployment_id}" + + +def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: + try: + call: Final = _CALL_KWARGS.validate_python(kwargs) + except ValidationError: + return None + litellm_params: Final = call.get("litellm_params") + metadata: Final = litellm_params.get("metadata") if litellm_params else None + model_info: Final = litellm_params.get("model_info") if litellm_params else None + model_group: Final = metadata.get("model_group") if metadata else None + deployment_id: Final = model_info.get("id") if model_info else None + if model_group is None or deployment_id is None: + return None + return model_group, str(deployment_id) + + +def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + return tuple( + _request_count_key(model_group, str(deployment["model_info"]["id"])) + for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + ) + + +def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: + return tuple(0 if value is None else int(value) for value in values) + + +def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: + values: Final = _MEMORY_COUNTS.validate_python(raw) + if values is None or len(values) != len(keys): + return (0,) * len(keys) + return _as_counts(values) + + +def _least_busy( + healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] +) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] + + +def _warn_unreadable(model_group: str, error: Exception) -> None: + verbose_router_logger.warning( + "least-busy routing could not read the shared in-flight counts for %s, " + "falling back to this worker's own counts: %s", + model_group, + error, + ) + + +def _warn_unwritable(key: str, error: Exception) -> None: + verbose_router_logger.warning("least-busy routing could not update the in-flight count under %s: %s", key, error) + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False @@ -20,195 +106,101 @@ class LeastBusyLoggingHandler(CustomLogger): def __init__(self, router_cache: DualCache): self.router_cache = router_cache + self.router_cache_id = str(id(router_cache)) - def log_pre_api_call(self, model, messages, kwargs): - """ - Log when a model is being used. + def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: + self._increment(kwargs, 1) - Caching based on model group. - """ - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - request_count_api_key: Final = f"{model_group}_request_count" - # update cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + 1 + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - except Exception: - pass + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - def log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - def _get_available_deployments( - self, - healthy_deployments: list, - all_deployments: dict, - ): - """ - Helper to get deployments using least busy strategy - """ - for d in healthy_deployments: - ## if healthy deployment not yet used - if d["model_info"]["id"] not in all_deployments: - all_deployments[d["model_info"]["id"]] = 0 - # map deployment to id - # pick least busy deployment - min_traffic = float("inf") - min_deployment = None - for k, v in all_deployments.items(): - if v < min_traffic: - min_traffic = v - min_deployment = k - if min_deployment is not None: - ## check if min deployment is a string, if so, cast it to int - for m in healthy_deployments: - if m["model_info"]["id"] == min_deployment: - return m - min_deployment = random.choice(healthy_deployments) - else: - min_deployment = random.choice(healthy_deployments) - return min_deployment + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 def get_available_deployments( - self, - model_group: str, - healthy_deployments: list, - ): - """ - Sync helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys))) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) - async def async_get_available_deployments(self, model_group: str, healthy_deployments: list): - """ - Async helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + async def async_get_available_deployments( + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys))) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) + + def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + local: Final = self.router_cache.increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local < 0: + self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + if redis_cache is None: + return + redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + except Exception as e: + _warn_unwritable(key, e) + + async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + local: Final = await self.router_cache.async_increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local is not None and local < 0: + await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + if redis_cache is None: + return + await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) + except Exception as e: + _warn_unwritable(key, e) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index bb6877a032b..805d4ff9080 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -32,6 +32,12 @@ def _average_latency(samples: Sequence[float]) -> float: return sum(samples) / len(samples) +def _ttft_seconds(elapsed: timedelta | float) -> float: + if isinstance(elapsed, timedelta): + return elapsed.total_seconds() + return float(elapsed) + + class LowestLatencyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 @@ -86,14 +92,13 @@ class LowestLatencyLoggingHandler(CustomLogger): # breaks JSON serialization when the router cache syncs to # Redis (issue #33169) response_ms = response_ms.total_seconds() - time_to_first_token_response_time = None + time_to_first_token: float | None = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time + time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time) final_value: float = response_ms - time_to_first_token: float | None = None total_tokens = 0 if isinstance(response_obj, ModelResponse): @@ -111,13 +116,6 @@ class LowestLatencyLoggingHandler(CustomLogger): else: final_value = response_seconds - if time_to_first_token_response_time is not None: - if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() - else: - ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) - # ------------ # Update usage # ------------ @@ -138,14 +136,14 @@ class LowestLatencyLoggingHandler(CustomLogger): ## Time to first token if time_to_first_token is not None: if ( - len(request_count_dict[id].get("time_to_first_token", [])) + len(request_count_dict[id].get("time_to_first_token_seconds", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ - 1: - ] + [time_to_first_token] + request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][ + "time_to_first_token_seconds" + ][1:] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -252,7 +250,7 @@ class LowestLatencyLoggingHandler(CustomLogger): {model_group}_map: { id: { "latency": [..] - "time_to_first_token": [..] + "time_to_first_token_seconds": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } @@ -273,14 +271,13 @@ class LowestLatencyLoggingHandler(CustomLogger): # breaks JSON serialization when the router cache syncs to # Redis (issue #33169) response_ms = response_ms.total_seconds() - time_to_first_token_response_time = None + time_to_first_token: float | None = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request - time_to_first_token_response_time = kwargs.get("completion_start_time", end_time) - start_time + time_to_first_token = _ttft_seconds(kwargs.get("completion_start_time", end_time) - start_time) final_value: float = response_ms total_tokens = 0 - time_to_first_token: float | None = None if isinstance(response_obj, ModelResponse): _usage: Final = getattr(response_obj, "usage", None) @@ -296,13 +293,6 @@ class LowestLatencyLoggingHandler(CustomLogger): final_value = float(normalized_value) else: final_value = response_seconds - - if time_to_first_token_response_time is not None: - if isinstance(time_to_first_token_response_time, timedelta): - ttft_seconds = time_to_first_token_response_time.total_seconds() - else: - ttft_seconds = time_to_first_token_response_time - time_to_first_token = safe_divide_seconds(ttft_seconds, completion_tokens) # ------------ # Update usage # ------------ @@ -328,14 +318,14 @@ class LowestLatencyLoggingHandler(CustomLogger): ## Time to first token if time_to_first_token is not None: if ( - len(request_count_dict[id].get("time_to_first_token", [])) + len(request_count_dict[id].get("time_to_first_token_seconds", [])) < self.routing_args.max_latency_list_size ): - request_count_dict[id].setdefault("time_to_first_token", []).append(time_to_first_token) + request_count_dict[id].setdefault("time_to_first_token_seconds", []).append(time_to_first_token) else: - request_count_dict[id]["time_to_first_token"] = request_count_dict[id]["time_to_first_token"][ - 1: - ] + [time_to_first_token] + request_count_dict[id]["time_to_first_token_seconds"] = request_count_dict[id][ + "time_to_first_token_seconds" + ][1:] + [time_to_first_token] if precise_minute not in request_count_dict[id]: request_count_dict[id][precise_minute] = {} @@ -433,7 +423,7 @@ class LowestLatencyLoggingHandler(CustomLogger): or float("inf") ) item_latency = item_map.get("latency", []) - item_ttft_latency = item_map.get("time_to_first_token", []) + item_ttft_latency = item_map.get("time_to_first_token_seconds", []) item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) diff --git a/litellm/router_strategy/simple_shuffle.py b/litellm/router_strategy/simple_shuffle.py index 65bcce0e532..860e89cea22 100644 --- a/litellm/router_strategy/simple_shuffle.py +++ b/litellm/router_strategy/simple_shuffle.py @@ -41,8 +41,7 @@ def simple_shuffle( ############## Check if 'weight' or 'rpm' or 'tpm' param set for a weighted pick ################# for weight_by in ["weight", "rpm", "tpm"]: - weight = healthy_deployments[0].get("litellm_params").get(weight_by, None) - if weight is not None: + if any(m["litellm_params"].get(weight_by) is not None for m in healthy_deployments): weights = [m["litellm_params"].get(weight_by, 0) for m in healthy_deployments] verbose_router_logger.debug("\nweight %s", weights) total_weight = sum(weights) diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py index 1e6412856a6..9699ab886b9 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", @@ -20,6 +20,7 @@ HEURISTIC_V1_TUNING_FIELDS: Final = ( "reasoning_override_min_score", "token_thresholds", "dimension_weights", + "custom_dimensions", "code_keywords", "reasoning_keywords", "technical_keywords", @@ -29,6 +30,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 +43,26 @@ 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, + exclude={ + "custom_dimensions": { + index: {"scoring_mode"} + for index, dimension in enumerate(validated.custom_dimensions) + if dimension.scoring_mode == "binary" + } + }, + ) return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index 9e7f457f631..ef29f7d8fd3 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -12,6 +12,7 @@ from typing_extensions import TypedDict from litellm import verbose_logger from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker if TYPE_CHECKING: @@ -36,10 +37,19 @@ _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0 class CooldownCache: - def __init__(self, cache: DualCache, default_cooldown_time: float): + def __init__( + self, + cache: DualCache, + default_cooldown_time: float, + redis_read_interval_seconds: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, + ): self.cache = cache self.default_cooldown_time = default_cooldown_time self.in_memory_cache = InMemoryCache() + self._cooldown_store = DualCache( + in_memory_cache=self.in_memory_cache, + default_redis_batch_cache_expiry=redis_read_interval_seconds, + ) # Initialize the masker with custom settings for exception strings self.exception_masker = SensitiveDataMasker( visible_prefix=50, # Show first 50 characters @@ -48,6 +58,21 @@ class CooldownCache: mask_short_values=False, # Truncate long messages only; keep short ones readable ) + @property + def cooldown_store(self) -> DualCache: + """ + The cache cooldown entries live in, with the router's Redis attached on first use. + + It is kept separate from the router-wide cache so that a key missing from memory is + re-read from Redis every `redis_read_interval_seconds` rather than on the router + cache's much longer batch interval, which is what lets a sibling replica see a + cooldown another replica wrote, and so that unrelated router keys cannot evict a + cooldown from the in-memory tier before it expires. Redis is attached lazily because + the router builds its cooldown cache before it wires up the shared Redis client. + """ + self._cooldown_store.attach_redis_cache(self.cache.redis_cache) + return self._cooldown_store + def _common_add_cooldown_logic( self, model_id: str, original_exception, exception_status, cooldown_time: float ) -> tuple[str, CooldownCacheValue]: @@ -93,7 +118,7 @@ class CooldownCache: ) # Set the cache with a TTL equal to the cooldown time - self.cache.set_cache( + self.cooldown_store.set_cache( value=cooldown_data, key=cooldown_key, ttl=_cooldown_time, @@ -122,13 +147,13 @@ class CooldownCache: cooldown_cache_value: Final = CooldownCacheValue(**result) # pyright: ignore[reportUnknownArgumentType] - result comes from an untyped cache read, not from our own code remaining: Final = (cooldown_cache_value["timestamp"] + cooldown_cache_value["cooldown_time"]) - current_time if remaining <= 0: - self.cache.in_memory_cache.delete_cache(key) + self.in_memory_cache.delete_cache(key) return None - current_expiry: Final = self.cache.in_memory_cache.ttl_dict.get(key) + current_expiry: Final = self.in_memory_cache.ttl_dict.get(key) if current_expiry is not None and current_expiry > current_time + remaining + 5: corrected_ttl: Final = min(remaining, _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS) - self.cache.in_memory_cache.delete_cache(key) - self.cache.in_memory_cache.set_cache(key, result, ttl=corrected_ttl) + self.in_memory_cache.delete_cache(key) + self.in_memory_cache.set_cache(key, result, ttl=corrected_ttl) return cooldown_cache_value async def async_get_active_cooldowns( @@ -137,12 +162,7 @@ class CooldownCache: # Generate the keys for the deployments keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] - # Retrieve the values for the keys using mget - ## more likely to be none if no models ratelimited. So just check redis every 1s - ## each redis call adds ~100ms latency. - - ## check in memory cache first - results: Final = await self.cache.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) + results: Final = await self.cooldown_store.async_batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = [] if results is None or all(v is None for v in results): @@ -164,7 +184,7 @@ class CooldownCache: # Generate the keys for the deployments keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids] # Retrieve the values for the keys using mget - results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] + results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] active_cooldowns: Final = [] current_time: Final = time.time() @@ -184,7 +204,7 @@ class CooldownCache: keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids] # Retrieve the values for the keys using mget - results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] + results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] min_cooldown_time: float | None = None # Process the results diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 4534fa114b3..f722b6fd20c 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache from litellm.constants import ( DEFAULT_COOLDOWN_TIME_SECONDS, DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS, @@ -558,9 +559,12 @@ def should_cooldown_based_on_allowed_fails_policy( When *allowed_fails_override* / *cooldown_time_override* are supplied they take precedence over the router-level values (used by deployment-level overrides). + The counter lives in the router's shared ``DualCache`` (Redis when configured), so + every worker process increments the same key and the threshold applies fleet-wide. + When *cache_key_suffix* is supplied the fail counter is keyed as - ``{deployment}:{cache_key_suffix}`` so that different exception types are - tracked independently per deployment. + ``deployment:{deployment}:allowed_fails:{cache_key_suffix}`` so that different + exception types are tracked independently per deployment. Returns: - True if fails exceed the allowed limit (should cooldown) @@ -584,16 +588,25 @@ def should_cooldown_based_on_allowed_fails_policy( else (litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS) ) - cache_key: Final = f"{deployment}:{cache_key_suffix}" if cache_key_suffix else deployment - current_fails: Final = litellm_router_instance.failed_calls.get_cache(key=cache_key) or 0 - updated_fails: Final = current_fails + 1 + base_key: Final = f"deployment:{deployment}:allowed_fails" + cache_key: Final = f"{base_key}:{cache_key_suffix}" if cache_key_suffix else base_key + updated_fails: Final = _increment_allowed_fails( + cache=litellm_router_instance.cache, cache_key=cache_key, ttl=cooldown_time + ) + return updated_fails > allowed_fails - if updated_fails > allowed_fails: - return True - else: - litellm_router_instance.failed_calls.set_cache(key=cache_key, value=updated_fails, ttl=cooldown_time) - return False +def _increment_allowed_fails(cache: DualCache, cache_key: str, ttl: float) -> int: + """ + Return the fleet-wide fail count. ``DualCache.increment_cache`` bumps the in-memory tier + before Redis and re-raises a Redis error, so a Redis outage degrades to this worker's own count. + """ + try: + return cache.increment_cache(key=cache_key, value=1, ttl=ttl) + except Exception as e: # noqa: BLE001 # a Redis outage must not stop failing deployments from cooling down + verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e) + local_fails: Final = cache.get_cache(key=cache_key, local_only=True) + return local_fails if isinstance(local_fails, int) else 0 def _is_allowed_fails_set_on_router( diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 0788c8db710..70362e60495 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -90,6 +90,7 @@ class PromptCachingDeploymentCheck(CustomLogger): enable_prompt_caching=( request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None ), + request_kwargs=request_kwargs, ) model_id_dict: Final = await prompt_cache.async_get_model_id( 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/agents.py b/litellm/types/agents.py index 2cb42ce3fac..dbaaab62d86 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,9 +1,9 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal from pydantic import BaseModel, PrivateAttr, StrictInt -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -172,6 +172,7 @@ class AugmentedAgentCard(AgentCard): class AgentObjectPermission(TypedDict, total=False): mcp_servers: list[str] | None mcp_access_groups: list[str] | None + mcp_toolsets: ReadOnly[Sequence[str] | None] mcp_tool_permissions: dict[str, list[str]] | None models: list[str] | None agents: list[str] | None 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/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 50c3515cf01..6306658ad0b 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -195,6 +195,10 @@ class AutoRouterBenchmarkTotals(BaseModel): avg_session_seconds: float avg_tokens_per_session: float spend: float = Field(description="What the routed traffic actually cost") + classifier_cost: float | None = Field( + description="Recorded LLM classifier cost already included in spend; null when any session turns predate " + "subtotal recording, and zero for an empty window" + ) saved_spend: float = Field( description="Signed dollars saved versus each router's savings baseline (derived from its hardest " "tier, or the configured override), from the same per-request savings record the usage tab reads" 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/utils.py b/litellm/types/utils.py index 010ff18d166..40a6d77482f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -580,6 +580,7 @@ CallTypesLiteral = Literal[ "search", "asearch", "_arealtime", + "_aresponses_websocket", "create_batch", "acreate_batch", "create_file", @@ -957,6 +958,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 +973,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 diff --git a/litellm/utils.py b/litellm/utils.py index 664d9597890..141ec323776 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -672,11 +672,19 @@ def load_credentials_from_list(kwargs: dict): CredentialAccessor: Final = getattr(sys.modules[__name__], "CredentialAccessor") credential_name: Final = kwargs.get("litellm_credential_name") - if credential_name and litellm.credential_list: - credential_accessor: Final[Mapping[str, object]] = CredentialAccessor.get_credential_values(credential_name) - for key, value in credential_accessor.items(): - if key not in kwargs: - kwargs[key] = value + if not credential_name: + return + credential: Final = CredentialAccessor.find_credential(credential_name) + if credential is None: + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + credential_name, + len(litellm.credential_list), + ) + return + for key, value in credential.credential_values.items(): + if key not in kwargs: + kwargs[key] = value def get_dynamic_callbacks( @@ -2805,6 +2813,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. @@ -5212,13 +5227,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: @@ -8687,6 +8705,8 @@ class ProviderConfigManager: return litellm.OpenRouterResponsesAPIConfig() elif litellm.LlmProviders.HOSTED_VLLM == provider: return litellm.HostedVLLMResponsesAPIConfig() + elif litellm.LlmProviders.FIREWORKS_AI == provider: + return litellm.FireworksAIResponsesAPIConfig() elif litellm.LlmProviders.BEDROCK_MANTLE == provider: # Both decisions are data-driven from the model's price-map entry, with # no model-name logic. Capability (can it serve Responses?) comes from diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 32563c32aa1..b1ffc1583e4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -30791,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", @@ -30808,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, @@ -37225,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, @@ -37265,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, @@ -37476,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, @@ -37516,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, diff --git a/pyproject.toml b/pyproject.toml index f4f238dd4b9..706b14b8a04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.101.0" +version = "1.102.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -67,7 +67,7 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.94", + "litellm-proxy-extras==0.4.95", "litellm-enterprise==0.1.65", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", @@ -328,7 +328,7 @@ members = ["enterprise", "litellm-proxy-extras"] profile = "black" [tool.commitizen] -version = "1.101.0" +version = "1.102.0" version_files = [ "pyproject.toml:^version", ] @@ -340,6 +340,7 @@ markers = [ "asyncio: mark test as an asyncio test", "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')", "no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests", + "requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension", ] filterwarnings = [ # Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc314..d63a69de76f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2956 + "limit": 2918 }, "ANN002": { "limit": 71 @@ -9,10 +9,10 @@ "limit": 806 }, "ANN201": { - "limit": 1979 + "limit": 1965 }, "ANN202": { - "limit": 831 + "limit": 829 }, "ANN204": { "limit": 683 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2916 + "limit": 2914 }, "C401": { "limit": 8 @@ -189,7 +189,7 @@ "limit": 0 }, "S110": { - "limit": 217 + "limit": 207 }, "S112": { "limit": 22 diff --git a/schema.prisma b/schema.prisma index 06ac177cca4..ccbab0fef10 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1509,6 +1509,8 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index adc4c0664be..485e118efd2 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -34,7 +34,7 @@ import subprocess import sys from pathlib import Path from types import MappingProxyType -from typing import NamedTuple +from typing import Final, NamedTuple if sys.version_info >= (3, 11): import tomllib @@ -42,7 +42,6 @@ else: import tomli as tomllib REPO_ROOT = Path(__file__).resolve().parent.parent -DEFAULT_BASE = "origin/litellm_internal_staging" DEFAULT_BUDGETS: tuple[str, ...] = ( "ruff-strict-budget.json", "type-discipline-budget.json", @@ -182,12 +181,15 @@ def regressions_for( def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("budgets", nargs="*", help="budget files to check") args = parser.parse_args() + from default_branch import resolve_base_ref + + base_ref: Final = resolve_base_ref(args.base, REPO_ROOT) budgets = args.budgets or list(DEFAULT_BUDGETS) - ref = _merge_base(args.base) + ref = _merge_base(base_ref) if not _ref_is_commit(ref): print( f"FAIL: base ref {ref!r} does not resolve to a commit, so the ratchet has nothing " @@ -204,14 +206,14 @@ def main() -> int: if base is None and head is None: continue if base is None: - print(f"skip {rel}: new file (no base at {args.base} to ratchet against)") + print(f"skip {rel}: new file (no base at {base_ref} to ratchet against)") continue checked.append(rel) regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel))) if regressions: print( - f"FAIL: budget limit(s) loosened vs base {args.base} (merge-base {ref[:12]}):" + f"FAIL: budget limit(s) loosened vs base {base_ref} (merge-base {ref[:12]}):" ) for reg in regressions: print(f" {reg.budget} {reg.rule}: {reg.detail}") @@ -223,7 +225,7 @@ def main() -> int: return 1 suffix = f" ({', '.join(checked)})" if checked else "" - print(f"OK: no budget limit increased vs base {args.base}{suffix}") + print(f"OK: no budget limit increased vs base {base_ref}{suffix}") return 0 diff --git a/scripts/default_branch.py b/scripts/default_branch.py new file mode 100644 index 00000000000..fb1852fd21c --- /dev/null +++ b/scripts/default_branch.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import argparse +import os +import subprocess +from pathlib import Path +from typing import Final + + +def _git(repo_root: Path, *args: str) -> str: + try: + result: Final = subprocess.run( + ["git", *args], + cwd=repo_root, + env={**os.environ, "GIT_TERMINAL_PROMPT": "0"}, + check=True, + capture_output=True, + text=True, + timeout=60, + ) + except (OSError, subprocess.SubprocessError) as exc: + raise SystemExit( + "Cannot verify the base branch against origin. Check remote access, " + "or supply an explicit base ref (--base / BASE_REF). " + f"Git operation failed: {exc}" + ) from exc + return result.stdout.strip() + + +def default_branch(repo_root: Path) -> str: + output: Final = _git(repo_root, "ls-remote", "--symref", "origin", "HEAD") + branches: Final = tuple( + line.removeprefix("ref: refs/heads/").removesuffix("\tHEAD") + for line in output.splitlines() + if line.startswith("ref: refs/heads/") and line.endswith("\tHEAD") + ) + if len(branches) != 1: + raise SystemExit("Origin did not advertise a default branch. Supply an explicit base ref (--base / BASE_REF).") + _git(repo_root, "check-ref-format", f"refs/heads/{branches[0]}") + return branches[0] + + +def resolve_base_ref(base_ref: str | None, repo_root: Path) -> str: + if base_ref: + return base_ref + branch: Final = default_branch(repo_root) + _git(repo_root, "fetch", "--quiet", "origin", f"+refs/heads/{branch}:refs/remotes/origin/{branch}") + return f"origin/{branch}" + + +def main() -> None: + parser: Final = argparse.ArgumentParser(description="Resolve the live default branch of origin.") + parser.add_argument("--base", help="Explicit comparison ref; skips default-branch discovery") + parser.add_argument("--repo-root", type=Path, default=Path.cwd()) + parser.add_argument("--branch", action="store_true", help="Print only the default branch name, without fetching") + args: Final = parser.parse_args() + print(default_branch(args.repo_root) if args.branch else resolve_base_ref(args.base, args.repo_root)) + + +if __name__ == "__main__": + main() diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index f245803408c..1abd415d237 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -7,7 +7,7 @@ # - anything staged -> scope is the staged files; changed-but-unstaged files # whose checks were skipped are called out # - nothing staged -> scope is the working tree's diff against the merge base -# with origin/litellm_internal_staging, untracked files included +# with origin's current default branch, untracked files included # The per-area checks: # - 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) @@ -33,8 +33,8 @@ set -eu # at a time instead of thrashing the machine. The wrapper exports # LITELLM_GATE_SLOT_HELD, so this re-exec happens exactly once and everything this # script spawns (make lint, the budget gates) skips its own acquisition. +script_dir=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0") if [ -z "${LITELLM_GATE_SLOT_HELD:-}" ]; then - script_dir=$(python3 -c 'import os, sys; print(os.path.dirname(os.path.realpath(sys.argv[1])))' "$0") exec python3 "$script_dir/gate_slot_lock.py" "$0" "$@" fi @@ -65,20 +65,24 @@ untracked=$(git ls-files --others --exclude-standard) if [ -n "$staged" ]; then scope=$staged else - git fetch --quiet origin litellm_internal_staging 2>/dev/null || true - merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || { - echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2 - echo " Fix: git fetch origin litellm_internal_staging" >&2 + base_ref=$(python3 "$script_dir/default_branch.py" --base "${BASE_REF:-}") || { + echo "check: FAIL" + exit 1 + } + export BASE_REF="$base_ref" + merge_base=$(git merge-base "$base_ref" HEAD 2>/dev/null) || { + echo "check: cannot resolve the merge base with $base_ref." >&2 + echo " Fix: fetch the base ref and provide BASE_REF=" >&2 echo "check: FAIL" exit 1 } scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u) if [ -z "$scope" ]; then - echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)" + echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs $base_ref)" echo "check: PASS" exit 0 fi - echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:" + echo "check: nothing staged; scoping to the working tree's diff against the merge base with $base_ref:" printf '%s\n' "$scope" | sed 's/^/ /' fi diff --git a/scripts/ruff_strict_gate.py b/scripts/ruff_strict_gate.py index bf070beeb0f..8da10dd76f0 100644 --- a/scripts/ruff_strict_gate.py +++ b/scripts/ruff_strict_gate.py @@ -24,7 +24,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent STRICT_CONFIG = REPO_ROOT / "ruff-strict.toml" BUDGET_PATH = REPO_ROOT / "ruff-strict-budget.json" TARGET = "litellm" -DEFAULT_BASE = "origin/litellm_internal_staging" _HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@") @@ -193,7 +192,7 @@ def ratcheted_budget(budget: dict, current: dict, base: dict) -> dict: } -def cmd_update(base_ref: str = DEFAULT_BASE) -> None: +def cmd_update(base_ref: str) -> None: """Ratchet each rule's limit down by the violations this branch fixed. The working-tree count is compared against a ruff pass over a detached @@ -212,13 +211,15 @@ def cmd_update(base_ref: str = DEFAULT_BASE) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("--update", action="store_true") args = parser.parse_args() + from default_branch import resolve_base_ref from gate_slot_lock import held_slot + base_ref: Final = resolve_base_ref(args.base, REPO_ROOT) with held_slot(): - cmd_update(args.base) if args.update else cmd_check(args.base) + cmd_update(base_ref) if args.update else cmd_check(base_ref) if __name__ == "__main__": diff --git a/scripts/test_quality_gate.py b/scripts/test_quality_gate.py index 4f4eeb17ec1..e486324c741 100644 --- a/scripts/test_quality_gate.py +++ b/scripts/test_quality_gate.py @@ -14,7 +14,7 @@ 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 +against the repository's default branch, 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. @@ -43,7 +43,6 @@ REPO_ROOT: Final = Path(__file__).resolve().parent.parent 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) @@ -240,7 +239,7 @@ def ratcheted_budget( }) -def cmd_update(base_ref: str = DEFAULT_BASE) -> None: +def cmd_update(base_ref: str) -> None: """Ratchet each rule's limit down by the violations this branch fixed.""" budget: Final = json.loads(BUDGET_PATH.read_text()) base_point: Final = resolve_base_point(base_ref) @@ -264,19 +263,20 @@ def cmd_seed() -> None: def main() -> None: parser: Final = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("--update", action="store_true") parser.add_argument("--seed", action="store_true") args: Final = parser.parse_args() + from default_branch import resolve_base_ref from gate_slot_lock import held_slot with held_slot(): if args.seed: cmd_seed() elif args.update: - cmd_update(args.base) + cmd_update(resolve_base_ref(args.base, REPO_ROOT)) else: - cmd_check(args.base) + cmd_check(resolve_base_ref(args.base, REPO_ROOT)) if __name__ == "__main__": diff --git a/scripts/type_check_gate.py b/scripts/type_check_gate.py index 763835e6d2e..78f74ec65a1 100644 --- a/scripts/type_check_gate.py +++ b/scripts/type_check_gate.py @@ -71,7 +71,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent BUDGET_PATH = REPO_ROOT / "basedpyright-code-budget.json" PYRIGHT_CONFIG = REPO_ROOT / "pyrightconfig.json" UV_LOCK = REPO_ROOT / "uv.lock" -DEFAULT_BASE = "origin/litellm_internal_staging" CACHE_FILE_PREFIX = "basedpyright-base-" CACHE_KEEP_ENTRIES = 8 ARTIFACT_NAME_PREFIX = "basedpyright-counts-" @@ -578,7 +577,7 @@ def ratcheted_budget( } -def cmd_update(current: Mapping[str, int], base_ref: str = DEFAULT_BASE) -> None: +def cmd_update(current: Mapping[str, int], base_ref: str) -> None: """Ratchet each rule's limit down by the errors this branch fixed. `current` is the working-tree count; the reference count comes @@ -666,12 +665,14 @@ def cmd_check(head: Mapping[str, int], base_ref: str) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("--update", action="store_true") parser.add_argument("--emit-counts-dir", type=Path) args = parser.parse_args() + from default_branch import resolve_base_ref from gate_slot_lock import held_slot + base_ref: Final = None if args.emit_counts_dir is not None else resolve_base_ref(args.base, REPO_ROOT) with held_slot(): ensure_typecheck_env() head = count_basedpyright(run_basedpyright()) @@ -679,10 +680,8 @@ def main() -> None: cmd_emit_counts( head, args.emit_counts_dir, _run(["git", "rev-parse", "HEAD"]).strip() ) - elif args.update: - cmd_update(head, args.base) - else: - cmd_check(head, args.base) + elif base_ref is not None: + cmd_update(head, base_ref) if args.update else cmd_check(head, base_ref) if __name__ == "__main__": diff --git a/scripts/type_discipline_gate.py b/scripts/type_discipline_gate.py index 5f6474f20bc..40e61cf7265 100644 --- a/scripts/type_discipline_gate.py +++ b/scripts/type_discipline_gate.py @@ -44,7 +44,6 @@ REPO_ROOT = Path(__file__).resolve().parent.parent CHECKER = REPO_ROOT / "scripts" / "check_type_discipline.py" BUDGET_PATH = REPO_ROOT / "type-discipline-budget.json" TARGET = "litellm" -DEFAULT_BASE = "origin/litellm_internal_staging" _HUNK = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@") _LINE = re.compile(r"^(?P.+?):(?P\d+): (?PLIT\d+) ") @@ -239,7 +238,7 @@ def _base_budget_rules(base_point: str) -> frozenset: return frozenset(json.loads(proc.stdout)) -def cmd_update(base_ref: str = DEFAULT_BASE) -> None: +def cmd_update(base_ref: str) -> None: """Ratchet each rule's limit down by the violations this branch fixed. The working-tree count is compared against a checker pass over a detached @@ -264,13 +263,15 @@ def cmd_update(base_ref: str = DEFAULT_BASE) -> None: def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--base", default=DEFAULT_BASE) + parser.add_argument("--base", help="Comparison ref (default: origin's current default branch)") parser.add_argument("--update", action="store_true") args = parser.parse_args() + from default_branch import resolve_base_ref from gate_slot_lock import held_slot + base_ref: Final = resolve_base_ref(args.base, REPO_ROOT) with held_slot(): - cmd_update(args.base) if args.update else cmd_check(args.base) + cmd_update(base_ref) if args.update else cmd_check(base_ref) if __name__ == "__main__": diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 389027bf5ca..923a4ca2da8 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -242,6 +242,22 @@ this with `litellm_license`. To tune the export cadence, set `LITELLM_BILLING_METRICS_EXPORT_INTERVAL_MS` through `gateway_extra_env` / `backend_extra_env` +### Prometheus metrics sidecar + +`gateway_metrics_port` adds a `metrics` sidecar +(`python -m litellm.proxy.prometheus_metrics_server`) to the gateway task that +aggregates the workers' samples over a shared task volume, so a scrape never +runs on an inference worker. The ALB never routes to that port and the tasks +security group only opens it to `gateway_metrics_scrape_cidrs`. Needs +`gateway_image` v1.101.0 or newer. See +[Prometheus metrics](https://docs.litellm.ai/docs/proxy/prometheus) for the +metrics themselves. + +```hcl +gateway_metrics_port = 4001 +gateway_metrics_scrape_cidrs = ["10.0.0.0/16"] +``` + ## Tenant deployment Every resource the stack creates is named `${tenant}-litellm-${env}` (or diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index 2fb1ca08205..aa2c3d558e2 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -212,6 +212,45 @@ locals { # pull the config from S3 first, so the command goes through `sh -c`; # otherwise we keep the image's ENTRYPOINT and only override `command`. gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}" + + metrics_enabled = var.gateway_metrics_port != null + metrics_multiproc_dir = "/tmp/litellm_prometheus_multiproc" + metrics_volume = "prometheus-multiproc" + metrics_env = local.metrics_enabled ? [{ name = "PROMETHEUS_MULTIPROC_DIR", value = local.metrics_multiproc_dir }] : [] + metrics_mount_points = local.metrics_enabled ? [{ sourceVolume = local.metrics_volume, containerPath = local.metrics_multiproc_dir }] : [] + metrics_health_cmd = "import socket; socket.create_connection(('127.0.0.1', ${coalesce(var.gateway_metrics_port, 0)}), timeout=2).close()" + + gateway_metrics_container = local.metrics_enabled ? [ + { + name = "metrics" + image = var.gateway_image + essential = false + entryPoint = ["python", "-m", "litellm.proxy.prometheus_metrics_server"] + command = ["--port", tostring(var.gateway_metrics_port)] + + portMappings = [{ containerPort = var.gateway_metrics_port, protocol = "tcp" }] + environment = local.metrics_env + mountPoints = local.metrics_mount_points + + healthCheck = { + command = ["CMD", "python", "-c", local.metrics_health_cmd] + interval = 30 + timeout = 5 + retries = 3 + startPeriod = 30 + } + + logConfiguration = { + logDriver = "awslogs" + options = { + awslogs-group = aws_cloudwatch_log_group.gateway.name + awslogs-region = var.region + awslogs-stream-prefix = "metrics" + } + } + } + ] : [] + backend_uvicorn_args = "--host 0.0.0.0 --port 4001" gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac" @@ -269,7 +308,7 @@ resource "aws_ecs_task_definition" "gateway" { execution_role_arn = aws_iam_role.task_execution.arn task_role_arn = aws_iam_role.task.arn - container_definitions = jsonencode([ + container_definitions = jsonencode(concat([ merge( { name = "gateway" @@ -283,8 +322,10 @@ resource "aws_ecs_task_definition" "gateway" { local.billing_metrics_env, local.gateway_extra_env_list, local.proxy_config_env, + local.metrics_env, ) - secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list) + secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list) + mountPoints = local.metrics_mount_points # Container-level healthCheck intentionally omitted — the wolfi # runtime image doesn't ship curl/wget. The ALB target group polls @@ -301,7 +342,14 @@ resource "aws_ecs_task_definition" "gateway" { }, local.gateway_proxy_overrides, ) - ]) + ], local.gateway_metrics_container)) + + dynamic "volume" { + for_each = local.metrics_enabled ? [1] : [] + content { + name = local.metrics_volume + } + } tags = local.tags } diff --git a/terraform/litellm/aws/examples/default/main.tf b/terraform/litellm/aws/examples/default/main.tf index 2eeaf6adb50..115003f7ddd 100644 --- a/terraform/litellm/aws/examples/default/main.tf +++ b/terraform/litellm/aws/examples/default/main.tf @@ -48,4 +48,7 @@ module "litellm" { backend_extra_env = var.backend_extra_env gateway_extra_secrets = var.gateway_extra_secrets backend_extra_secrets = var.backend_extra_secrets + + gateway_metrics_port = var.gateway_metrics_port + gateway_metrics_scrape_cidrs = var.gateway_metrics_scrape_cidrs } diff --git a/terraform/litellm/aws/examples/default/terraform.tfvars.example b/terraform/litellm/aws/examples/default/terraform.tfvars.example index 59301ea6aa5..880ecf56555 100644 --- a/terraform/litellm/aws/examples/default/terraform.tfvars.example +++ b/terraform/litellm/aws/examples/default/terraform.tfvars.example @@ -102,6 +102,13 @@ env = "stage" # } # } +# ---------- Prometheus metrics sidecar ---------- +# Serve /metrics from a sidecar in the gateway task instead of the inference +# workers. The port is not behind the ALB and has no auth: open it only to +# your Prometheus subnets. +# gateway_metrics_port = 4001 +# gateway_metrics_scrape_cidrs = ["10.0.0.0/16"] + # ---------- Extra env / secrets ---------- # Plain-text env vars (non-sensitive). Land directly in the ECS task def. # gateway_extra_env = { diff --git a/terraform/litellm/aws/examples/default/variables.tf b/terraform/litellm/aws/examples/default/variables.tf index d8ab56b13af..f8140266fca 100644 --- a/terraform/litellm/aws/examples/default/variables.tf +++ b/terraform/litellm/aws/examples/default/variables.tf @@ -158,3 +158,15 @@ variable "backend_extra_secrets" { type = map(string) default = {} } + +variable "gateway_metrics_port" { + description = "Port for the Prometheus metrics sidecar in the gateway task. Null keeps /metrics on the gateway port only." + type = number + default = null +} + +variable "gateway_metrics_scrape_cidrs" { + description = "CIDRs allowed to scrape gateway_metrics_port." + type = list(string) + default = [] +} diff --git a/terraform/litellm/aws/network.tf b/terraform/litellm/aws/network.tf index 4563eefbba5..c54949b4a59 100644 --- a/terraform/litellm/aws/network.tf +++ b/terraform/litellm/aws/network.tf @@ -156,6 +156,17 @@ resource "aws_security_group" "tasks" { security_groups = [aws_security_group.alb.id] } + dynamic "ingress" { + for_each = local.metrics_enabled && length(var.gateway_metrics_scrape_cidrs) > 0 ? [1] : [] + content { + description = "Prometheus scrapers to the gateway metrics sidecar" + from_port = var.gateway_metrics_port + to_port = var.gateway_metrics_port + protocol = "tcp" + cidr_blocks = var.gateway_metrics_scrape_cidrs + } + } + egress { description = "All egress (LLM providers, RDS, Redis)" from_port = 0 diff --git a/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl b/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl new file mode 100644 index 00000000000..de839d7e5be --- /dev/null +++ b/terraform/litellm/aws/tests/metrics_sidecar.tftest.hcl @@ -0,0 +1,108 @@ +# Plan-only coverage for the Prometheus metrics sidecar wiring. Offline via +# mock_provider, same as byo_infrastructure.tftest.hcl. + +mock_provider "aws" { + mock_data "aws_iam_policy_document" { + defaults = { + json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}" + } + } +} +mock_provider "random" {} + +variables { + region = "us-east-1" + tenant = "acme" + env = "test" + allow_plaintext_alb = true + azs = ["us-east-1a", "us-east-1b"] +} + +run "defaults_change_nothing" { + command = plan + + assert { + condition = alltrue([ + length(local.gateway_metrics_container) == 0, + length(local.metrics_env) == 0, + length(local.metrics_mount_points) == 0, + length([for r in aws_security_group.tasks.ingress : r if r.description == "Prometheus scrapers to the gateway metrics sidecar"]) == 0, + ]) + error_message = "The metrics sidecar, its env, its volume, and its security-group rule must all be absent by default." + } +} + +run "metrics_port_adds_a_sidecar_volume_and_scrape_rule" { + command = plan + + variables { + gateway_metrics_port = 9464 + gateway_metrics_scrape_cidrs = ["10.20.0.0/16"] + } + + assert { + condition = length(local.metrics_env) == 1 && local.metrics_env[0].name == "PROMETHEUS_MULTIPROC_DIR" && local.metrics_env[0].value == "/tmp/litellm_prometheus_multiproc" + error_message = "The gateway workers must write multiprocess samples to the shared dir." + } + + assert { + condition = length(local.metrics_mount_points) == 1 && local.metrics_mount_points[0].sourceVolume == "prometheus-multiproc" && local.metrics_mount_points[0].containerPath == "/tmp/litellm_prometheus_multiproc" + error_message = "Gateway and sidecar must mount the same task volume at the multiproc dir." + } + + assert { + condition = alltrue([ + length(local.gateway_metrics_container) == 1, + local.gateway_metrics_container[0].name == "metrics", + local.gateway_metrics_container[0].essential == false, + join(" ", local.gateway_metrics_container[0].entryPoint) == "python -m litellm.proxy.prometheus_metrics_server", + join(" ", local.gateway_metrics_container[0].command) == "--port 9464", + one(local.gateway_metrics_container[0].portMappings).containerPort == 9464, + one(local.gateway_metrics_container[0].environment).value == "/tmp/litellm_prometheus_multiproc", + one(local.gateway_metrics_container[0].mountPoints).sourceVolume == "prometheus-multiproc", + strcontains(local.gateway_metrics_container[0].healthCheck.command[3], "9464"), + ]) + error_message = "The metrics sidecar must run prometheus_metrics_server on the configured port, share the multiproc volume, and health-check that port." + } + + assert { + condition = length(aws_ecs_task_definition.gateway.volume) == 1 && one(aws_ecs_task_definition.gateway.volume).name == "prometheus-multiproc" + error_message = "The gateway task must declare the multiproc volume." + } + + assert { + condition = length([ + for r in aws_security_group.tasks.ingress : r + if r.from_port == 9464 && r.to_port == 9464 && r.protocol == "tcp" && r.cidr_blocks == tolist(["10.20.0.0/16"]) + ]) == 1 + error_message = "The scrape CIDRs must be allowed to reach the metrics port on the tasks security group." + } + + assert { + condition = aws_lb_target_group.gateway.port == 4000 && one(aws_ecs_service.gateway.load_balancer).container_port == 4000 + error_message = "The ALB must keep targeting the gateway port only; the metrics port is never load balanced." + } +} + +run "metrics_port_without_scrape_cidrs_opens_nothing" { + command = plan + + variables { + gateway_metrics_port = 9464 + } + + assert { + condition = length(local.gateway_metrics_container) == 1 && length([for r in aws_security_group.tasks.ingress : r if r.from_port == 9464]) == 0 + error_message = "Without scrape CIDRs the sidecar runs but the metrics port stays closed to everything but the ALB group." + } +} + +run "metrics_port_may_not_reuse_the_gateway_port" { + command = plan + + variables { + gateway_metrics_port = 4000 + } + + expect_failures = [var.gateway_metrics_port] +} diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 522138953d6..667f6db63c9 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -549,6 +549,44 @@ variable "proxy_config" { default = {} } +# ---------- Prometheus metrics sidecar ---------- + +variable "gateway_metrics_port" { + description = <<-EOT + Serve Prometheus /metrics from a `metrics` sidecar container in the + gateway task on this port (1-65535, not 4000), so a scrape never runs on + an inference worker. The sidecar runs the gateway image with + `python -m litellm.proxy.prometheus_metrics_server` and aggregates the + workers' PROMETHEUS_MULTIPROC_DIR samples over a task volume. Null (the + default) leaves /metrics on the gateway port only. The sidecar port has + no virtual-key auth and is not routed through the ALB; open it to your + scrapers with gateway_metrics_scrape_cidrs. Needs gateway_image v1.101.0 + or newer. + EOT + type = number + default = null + + validation { + condition = var.gateway_metrics_port == null || (var.gateway_metrics_port >= 1 && var.gateway_metrics_port <= 65535 && var.gateway_metrics_port != 4000) + error_message = "gateway_metrics_port must be between 1 and 65535 and must not be 4000 (the gateway port)." + } +} + +variable "gateway_metrics_scrape_cidrs" { + description = <<-EOT + CIDR blocks allowed to reach gateway_metrics_port on the gateway tasks + (your Prometheus or collector subnets). Empty by default, so only the + ALB can reach the tasks. Ignored when gateway_metrics_port is null. + EOT + type = list(string) + default = [] + + validation { + condition = alltrue([for c in var.gateway_metrics_scrape_cidrs : can(cidrnetmask(c))]) + error_message = "gateway_metrics_scrape_cidrs must contain valid IPv4 CIDR blocks." + } +} + variable "log_retention_days" { description = "CloudWatch log retention for the three services." type = number diff --git a/terraform/provider/RELEASING.md b/terraform/provider/RELEASING.md index 59f4c5f066c..5b72621c291 100644 --- a/terraform/provider/RELEASING.md +++ b/terraform/provider/RELEASING.md @@ -79,7 +79,7 @@ Before publishing to the Terraform Registry: ## What a change needs -1. **Land it in `BerriAI/litellm`.** Open a PR against `litellm_internal_staging` with the source change and a `CHANGELOG.md` entry under `[Unreleased]`. CI runs `gofmt`, `go vet`, build, tests and the endpoint-drift audit. A change that breaks existing configurations or state must say so in the changelog: the version number cannot signal it any more +1. **Land it in `BerriAI/litellm`.** Open a PR against the repository's current default branch with the source change and a `CHANGELOG.md` entry under `[Unreleased]`. CI runs `gofmt`, `go vet`, build, tests and the endpoint-drift audit. A change that breaks existing configurations or state must say so in the changelog: the version number cannot signal it any more 2. **Wait for the next LiteLLM release.** The nightly dev release carries it within a day; it reaches a stable version on the next stable cut 3. **Verify** (optional): the version appears at https://registry.terraform.io/providers/BerriAI/litellm and https://github.com/BerriAI/terraform-provider-litellm/releases. If the tag is on the mirror but there is no release, the goreleaser run failed: https://github.com/BerriAI/terraform-provider-litellm/actions 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/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/models.py b/tests/e2e/models.py index 0b718ae94f3..b4fd1d0ac1b 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -655,6 +655,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 @@ -671,6 +692,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]]): 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/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/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/llm_translation/test_bedrock_embedding.py b/tests/llm_translation/test_bedrock_embedding.py index 56baed141da..1fc05b43b23 100644 --- a/tests/llm_translation/test_bedrock_embedding.py +++ b/tests/llm_translation/test_bedrock_embedding.py @@ -1,14 +1,12 @@ import json import os -from datetime import datetime -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import Mock, patch import pytest import base64 -import httpx import litellm -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import HTTPHandler titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} @@ -394,8 +392,6 @@ def test_bedrock_embedding_uses_correct_region_when_specified(): os.environ["AWS_REGION_NAME"] = original_region_name else: os.environ.pop("AWS_REGION_NAME", None) - - def test_bedrock_embedding_region_bug_reproduction(): """ Reproduces the bug where aws_region_name is ignored when passed explicitly. @@ -458,13 +454,3 @@ def test_bedrock_embedding_region_bug_reproduction(): os.environ["AWS_REGION_NAME"] = original_region_name else: os.environ.pop("AWS_REGION_NAME", None) - - -def test_bedrock_titan_g1_text_02_model_info(): - """Test that amazon.titan-embed-g1-text-02 has correct pricing metadata""" - model_info = litellm.get_model_info("amazon.titan-embed-g1-text-02") - assert model_info is not None, "Model info should not be None" - assert model_info["litellm_provider"] == "bedrock" - assert model_info["mode"] == "embedding" - assert model_info["input_cost_per_token"] == 1e-07 - assert model_info["max_input_tokens"] == 8192 diff --git a/tests/llm_translation/test_bedrock_embedding_pricing.py b/tests/llm_translation/test_bedrock_embedding_pricing.py deleted file mode 100644 index 099d73fed87..00000000000 --- a/tests/llm_translation/test_bedrock_embedding_pricing.py +++ /dev/null @@ -1,34 +0,0 @@ -""" -Tests for AWS Bedrock embedding model pricing in the model cost map. - -Regression test for the Amazon Titan Text Embeddings V2 commercial price, -which was previously set 10x too high (2e-07 instead of 2e-08). -AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens -(= $0.00002 per 1K tokens = 2e-08 per token). -""" - -import importlib - - -class TestBedrockEmbeddingPricing: - """Test suite for Bedrock embedding model pricing in the cost map.""" - - def test_titan_embed_v2_commercial_input_cost(self, monkeypatch): - """Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08).""" - # Scope the local-cost-map flag to this test only, so it does not leak - # into sibling tests. monkeypatch restores the environment on teardown. - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - - import litellm.litellm_core_utils.get_model_cost_map - import litellm - - # Reload so the cost map is re-read from the local file with the flag set. - importlib.reload(litellm.litellm_core_utils.get_model_cost_map) - importlib.reload(litellm) - - model = litellm.model_cost["amazon.titan-embed-text-v2:0"] - - assert model["input_cost_per_token"] == 2e-08 - assert model["output_cost_per_token"] == 0.0 - assert model["litellm_provider"] == "bedrock" - assert model["mode"] == "embedding" diff --git a/tests/llm_translation/test_bedrock_govcloud.py b/tests/llm_translation/test_bedrock_govcloud.py index e69a95c714d..3ac1fa7cf2e 100644 --- a/tests/llm_translation/test_bedrock_govcloud.py +++ b/tests/llm_translation/test_bedrock_govcloud.py @@ -40,38 +40,6 @@ class TestBedrockGovCloudSupport: assert "us-gov-east-1" in all_regions assert "us-gov-west-1" in all_regions - def test_govcloud_models_in_model_cost(self): - """Test that GovCloud models are present in model cost configuration""" - from litellm import model_cost - - # Test Claude models in GovCloud - assert ( - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - in model_cost - ) - assert ( - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0" - in model_cost - ) - assert ( - "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost - ) - assert ( - "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" in model_cost - ) - assert "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0" in model_cost - assert "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0" in model_cost - - # Test Llama models in GovCloud - assert "bedrock/us-gov-east-1/meta.llama3-8b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-east-1/meta.llama3-70b-instruct-v1:0" in model_cost - assert "bedrock/us-gov-west-1/meta.llama3-70b-instruct-v1:0" in model_cost - - # Test Titan models in GovCloud - assert "bedrock/us-gov-east-1/amazon.titan-text-lite-v1" in model_cost - assert "bedrock/us-gov-west-1/amazon.titan-text-lite-v1" in model_cost - def test_govcloud_model_routing(self): """Test that GovCloud models are routed correctly""" # Test Claude model routing @@ -148,135 +116,6 @@ class TestBedrockGovCloudSupport: assert not any("us-gov-east-1" in model for model in litellm.bedrock_models) assert not any("us-gov-west-1" in model for model in litellm.bedrock_models) - def test_govcloud_model_cost_properties(self): - """Test that GovCloud models have proper cost configuration""" - from litellm import model_cost - - # Check a specific GovCloud model has all required properties - govcloud_model = model_cost[ - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ] - - assert "max_tokens" in govcloud_model - assert "max_input_tokens" in govcloud_model - assert "max_output_tokens" in govcloud_model - assert "input_cost_per_token" in govcloud_model - assert "output_cost_per_token" in govcloud_model - assert govcloud_model["litellm_provider"] == "bedrock" - assert govcloud_model["mode"] == "chat" - - def test_govcloud_model_pricing_verification(self): - """Test that GovCloud models have correct pricing that differs from base models""" - from litellm import model_cost - - # Claude Haiku 4.5 commercial list pricing is under the us.* inference profile id - base_model = "us.anthropic.claude-haiku-4-5-20251001-v1:0" - gov_east_model = ( - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ) - gov_west_model = ( - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0" - ) - - # Verify base model pricing (us.* inference profile: $1.10/$5.50 per MTok) - base_pricing = model_cost[base_model] - assert base_pricing["input_cost_per_token"] == 1.1e-06 - assert base_pricing["output_cost_per_token"] == 5.5e-06 - - # Verify GovCloud models have different (higher) pricing - gov_east_pricing = model_cost[gov_east_model] - gov_west_pricing = model_cost[gov_west_model] - - # GovCloud models should have ~20% higher pricing than base models - assert gov_east_pricing["input_cost_per_token"] == 1.2e-06 - assert gov_east_pricing["output_cost_per_token"] == 6e-06 - assert gov_west_pricing["input_cost_per_token"] == 1.2e-06 - assert gov_west_pricing["output_cost_per_token"] == 6e-06 - - # Verify the pricing difference is approximately 20% - assert ( - abs( - gov_east_pricing["input_cost_per_token"] - / base_pricing["input_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_east_pricing["output_cost_per_token"] - / base_pricing["output_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_west_pricing["input_cost_per_token"] - / base_pricing["input_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - assert ( - abs( - gov_west_pricing["output_cost_per_token"] - / base_pricing["output_cost_per_token"] - - 1.2 - ) - < 0.15 - ) - - # Test Claude 3 Haiku pricing - base_haiku_model = "anthropic.claude-3-haiku-20240307-v1:0" - gov_east_haiku_model = ( - "bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0" - ) - gov_west_haiku_model = ( - "bedrock/us-gov-west-1/anthropic.claude-3-haiku-20240307-v1:0" - ) - - # Verify base Haiku model pricing - base_haiku_pricing = model_cost[base_haiku_model] - assert base_haiku_pricing["input_cost_per_token"] == 2.5e-07 # 0.00000025 - assert base_haiku_pricing["output_cost_per_token"] == 1.25e-06 # 0.00000125 - - # Verify GovCloud Haiku models have different (higher) pricing - gov_east_haiku_pricing = model_cost[gov_east_haiku_model] - gov_west_haiku_pricing = model_cost[gov_west_haiku_model] - - # GovCloud Haiku models should have 20% higher pricing than base models - assert ( - gov_east_haiku_pricing["input_cost_per_token"] == 3e-07 - ) # 0.0000003 (20% higher) - assert ( - gov_east_haiku_pricing["output_cost_per_token"] == 1.5e-06 - ) # 0.0000015 (20% higher) - assert ( - gov_west_haiku_pricing["input_cost_per_token"] == 3e-07 - ) # 0.0000003 (20% higher) - assert ( - gov_west_haiku_pricing["output_cost_per_token"] == 1.5e-06 - ) # 0.0000015 (20% higher) - - # Verify the pricing difference is exactly 20% - assert ( - gov_east_haiku_pricing["input_cost_per_token"] - == base_haiku_pricing["input_cost_per_token"] * 1.2 - ) - assert ( - gov_east_haiku_pricing["output_cost_per_token"] - == base_haiku_pricing["output_cost_per_token"] * 1.2 - ) - assert ( - gov_west_haiku_pricing["input_cost_per_token"] - == base_haiku_pricing["input_cost_per_token"] * 1.2 - ) - assert ( - gov_west_haiku_pricing["output_cost_per_token"] - == base_haiku_pricing["output_cost_per_token"] * 1.2 - ) - @patch("litellm.completion") def test_govcloud_completion_cost_calculation(self, mock_completion): """Test that completion requests use correct pricing for GovCloud models""" diff --git a/tests/llm_translation/test_crusoe.py b/tests/llm_translation/test_crusoe.py index 56aa4e4cd42..576428684fc 100644 --- a/tests/llm_translation/test_crusoe.py +++ b/tests/llm_translation/test_crusoe.py @@ -4,7 +4,6 @@ Tests for Crusoe provider integration import os from unittest import mock -import litellm CRUSOE_API_BASE = "https://managed-inference-api-proxy.crusoecloud.com/v1" @@ -71,38 +70,3 @@ def test_get_llm_provider_crusoe(): ) assert model == "meta-llama/Llama-3.3-70B-Instruct" assert provider == "crusoe" - - -def test_crusoe_models_configuration(): - """Test that Crusoe models are configured correctly""" - from litellm import get_model_info - - original_model_cost = litellm.model_cost - original_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") - try: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - crusoe_models = [ - "crusoe/meta-llama/Llama-3.3-70B-Instruct", - "crusoe/deepseek-ai/DeepSeek-R1-0528", - "crusoe/deepseek-ai/DeepSeek-V3-0324", - "crusoe/Qwen/Qwen3-235B-A22B-Instruct-2507", - "crusoe/moonshotai/Kimi-K2-Thinking", - "crusoe/openai/gpt-oss-120b", - "crusoe/google/gemma-3-12b-it", - ] - - for model in crusoe_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - assert model_info.get("litellm_provider") == "crusoe", ( - f"{model} should have crusoe as provider" - ) - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - finally: - litellm.model_cost = original_model_cost - if original_env is None: - os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None) - else: - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = original_env diff --git a/tests/llm_translation/test_hyperbolic.py b/tests/llm_translation/test_hyperbolic.py index 78817fbd902..b7206e40a4e 100644 --- a/tests/llm_translation/test_hyperbolic.py +++ b/tests/llm_translation/test_hyperbolic.py @@ -1,8 +1,3 @@ -import os -from datetime import datetime -from unittest.mock import MagicMock - -import pytest import litellm @@ -69,34 +64,6 @@ def test_hyperbolic_in_provider_lists(): assert "https://api.hyperbolic.xyz/v1" in openai_compatible_endpoints -def test_hyperbolic_models_configuration(): - """Test that Hyperbolic models are properly configured""" - import json - - # Load model configuration directly from the JSON file - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path, "r") as f: - model_data = json.load(f) - - # Test a few key models - test_models = [ - "hyperbolic/deepseek-ai/DeepSeek-V3", - "hyperbolic/Qwen/Qwen2.5-Coder-32B-Instruct", - "hyperbolic/deepseek-ai/DeepSeek-R1", - ] - - for model in test_models: - assert model in model_data - model_info = model_data[model] - assert model_info["litellm_provider"] == "hyperbolic" - assert model_info["mode"] == "chat" - assert "max_tokens" in model_info - assert "input_cost_per_token" in model_info - assert "output_cost_per_token" in model_info - - def test_hyperbolic_supported_params(): """Test that supported OpenAI parameters are correctly configured""" from litellm.llms.hyperbolic.chat.transformation import HyperbolicChatConfig diff --git a/tests/llm_translation/test_lambda_ai.py b/tests/llm_translation/test_lambda_ai.py index 7ae18828d3f..edba459b352 100644 --- a/tests/llm_translation/test_lambda_ai.py +++ b/tests/llm_translation/test_lambda_ai.py @@ -8,7 +8,6 @@ from unittest import mock import pytest import litellm -from litellm import completion from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig @@ -103,48 +102,6 @@ async def test_lambda_ai_completion_call(): raise -def test_lambda_ai_models_configuration(): - """Test that Lambda AI models are configured correctly""" - from litellm import get_model_info - - # Reload model cost map to pick up local changes - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - # Clear and repopulate lambda_ai_models list after reloading model_cost - litellm.lambda_ai_models = set() - litellm.add_known_models() - - # Some Lambda AI models to test - lambda_ai_models = [ - "lambda_ai/deepseek-llama3.3-70b", - "lambda_ai/hermes3-8b", - "lambda_ai/llama3.1-8b-instruct", - "lambda_ai/llama3.2-11b-vision-instruct", - "lambda_ai/qwen25-coder-32b-instruct", - ] - - for model in lambda_ai_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - assert ( - model_info.get("litellm_provider") == "lambda_ai" - ), f"{model} should have lambda_ai as provider" - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - assert ( - model_info.get("supports_function_calling") is True - ), f"{model} should support function calling" - assert ( - model_info.get("supports_system_messages") is True - ), f"{model} should support system messages" - - # Check vision support for vision models - if "vision" in model: - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - def test_lambda_ai_model_list_populated(): """Test that lambda_ai_models list is populated correctly""" # Ensure we're using local model cost map and repopulate models diff --git a/tests/llm_translation/test_morph.py b/tests/llm_translation/test_morph.py index b91d1810d38..752fb3b9083 100644 --- a/tests/llm_translation/test_morph.py +++ b/tests/llm_translation/test_morph.py @@ -68,24 +68,6 @@ def test_morph_in_provider_lists(): ) -def test_morph_model_info(): - """Test that morph models have correct configuration.""" - import litellm - - model_info = litellm.get_model_info("morph/morph-v3-large") - - assert model_info["litellm_provider"] == "morph" - assert model_info["mode"] == "chat" - assert model_info["max_tokens"] == 16000 - assert model_info["max_input_tokens"] == 16000 - assert model_info["max_output_tokens"] == 16000 - assert model_info["input_cost_per_token"] == 9e-07 # $0.9/1M tokens - assert model_info["output_cost_per_token"] == 1.9e-06 # $1.9/1M tokens - assert model_info["supports_function_calling"] is False - assert model_info["supports_vision"] is False - assert model_info["supports_system_messages"] is True - - def test_morph_supported_params(): """Test that MorphChatConfig returns correct supported parameters.""" config = MorphChatConfig() diff --git a/tests/llm_translation/test_openai_o1.py b/tests/llm_translation/test_openai_o1.py index e188a3af647..fd25e04d67d 100644 --- a/tests/llm_translation/test_openai_o1.py +++ b/tests/llm_translation/test_openai_o1.py @@ -1,15 +1,11 @@ -import json import os -from datetime import datetime -from unittest.mock import AsyncMock, patch, MagicMock +from unittest.mock import patch - -import httpx import pytest import litellm -from litellm import Choices, Message, ModelResponse +from litellm import ModelResponse from base_llm_unit_tests import BaseLLMChatTest, BaseOSeriesModelsTest @@ -74,7 +70,6 @@ async def test_o1_handle_tool_calling_optional_params( - max_tokens is translated to 'max_completion_tokens' - role 'system' is translated to 'user' """ - from openai import AsyncOpenAI from litellm.utils import ProviderConfigManager from litellm.types.utils import LlmProviders @@ -186,15 +181,6 @@ class TestOpenAIO3(BaseOSeriesModelsTest, BaseLLMChatTest): pass -def test_o1_supports_vision(): - """Test that o1 supports vision""" - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - for k, v in litellm.model_cost.items(): - if k.startswith("o1") and v.get("litellm_provider") == "openai": - assert v.get("supports_vision") is True, f"{k} does not support vision" - - def test_o3_reasoning_effort(): resp = litellm.completion( model="o3-mini", diff --git a/tests/llm_translation/test_v0.py b/tests/llm_translation/test_v0.py index 95708dd855a..e96022e1e22 100644 --- a/tests/llm_translation/test_v0.py +++ b/tests/llm_translation/test_v0.py @@ -8,7 +8,6 @@ from unittest import mock import pytest import litellm -from litellm import completion from litellm.llms.v0.chat.transformation import V0ChatConfig @@ -111,33 +110,3 @@ def test_v0_supported_params(): ] assert set(supported_params) == set(expected_params) - - -def test_v0_models_configuration(): - """Test that v0 models are configured correctly""" - from litellm import get_model_info - - # Reload model cost map to pick up local changes - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - # All v0 models - v0_models = ["v0/v0-1.0-md", "v0/v0-1.5-md", "v0/v0-1.5-lg"] - - for model in v0_models: - model_info = get_model_info(model) - assert model_info is not None, f"Model info not found for {model}" - # All v0 models support vision (multimodal) - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - assert ( - model_info.get("litellm_provider") == "v0" - ), f"{model} should have v0 as provider" - assert model_info.get("mode") == "chat", f"{model} should be in chat mode" - assert ( - model_info.get("supports_function_calling") is True - ), f"{model} should support function calling" - assert ( - model_info.get("supports_system_messages") is True - ), f"{model} should support system messages" diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index 2de83778f1c..38ccfd91f95 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -1,8 +1,6 @@ # What is this? ## Unit testing for the 'get_model_info()' function import os -import traceback -import json from typing import List, Dict, Any @@ -11,7 +9,7 @@ import pytest import litellm from litellm import get_model_info -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch def test_get_model_info_simple_model_name(): @@ -49,34 +47,6 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch): assert model_info["input_cost_per_token"] == 0.0 -def test_get_model_info_shows_correct_supports_vision(): - info = litellm.get_model_info("gemini/gemini-2.0-flash") - print("info", info) - assert info["supports_vision"] is True - - -def test_get_model_info_shows_assistant_prefill(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - info = litellm.get_model_info("deepseek/deepseek-chat") - print("info", info) - assert info.get("supports_assistant_prefill") is True - - -def test_get_model_info_shows_supports_prompt_caching(): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - info = litellm.get_model_info("deepseek/deepseek-chat") - print("info", info) - assert info.get("supports_prompt_caching") is True - - -def test_get_model_info_finetuned_models(): - info = litellm.get_model_info("ft:gpt-3.5-turbo:my-org:custom_suffix:id") - print("info", info) - assert info["input_cost_per_token"] == 0.000003 - - def test_get_model_info_gemini_pro(): info = litellm.get_model_info("gemini-2.0-flash") print("info", info) @@ -219,7 +189,7 @@ def test_model_info_bedrock_converse_enforcement(monkeypatch): def test_get_model_info_custom_provider(): # Custom provider example copied from https://docs.litellm.ai/docs/providers/custom_llm_server: import litellm - from litellm import CustomLLM, completion, get_llm_provider + from litellm import CustomLLM, completion class MyCustomLLM(CustomLLM): def completion(self, *args, **kwargs) -> litellm.ModelResponse: diff --git a/tests/local_testing/test_least_busy_routing.py b/tests/local_testing/test_least_busy_routing.py index 18ab8bf779d..fb83d4e601f 100644 --- a/tests/local_testing/test_least_busy_routing.py +++ b/tests/local_testing/test_least_busy_routing.py @@ -33,8 +33,8 @@ def test_model_added(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"gpt-3.5-turbo_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = "gpt-3.5-turbo_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 def test_get_available_deployments(): @@ -52,8 +52,8 @@ def test_get_available_deployments(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"{model_group}_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = f"{model_group}_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 # test_get_available_deployments() @@ -104,15 +104,20 @@ async def test_router_get_available_deployments(async_test): router.leastbusy_logger.test_flag = True model_group = "azure-model" - request_count_dict = {1: 10, 2: 54, 3: 100} - cache_key = f"{model_group}_request_count" + request_count_dict = {"1": 10, "2": 54, "3": 100} + cache_keys = { + deployment_id: f"{model_group}_request_count:{deployment_id}" + for deployment_id in request_count_dict + } if async_test is True: - await router.cache.async_set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + await router.cache.async_set_cache(key=cache_keys[deployment_id], value=count) deployment = await router.async_get_available_deployment( model=model_group, messages=None, request_kwargs={} ) else: - router.cache.set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + router.cache.set_cache(key=cache_keys[deployment_id], value=count) deployment = router.get_available_deployment(model=model_group, messages=None) print(f"deployment: {deployment}") assert deployment["model_info"]["id"] == "1" @@ -124,15 +129,18 @@ async def test_router_get_available_deployments(async_test): messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - return_dict = router.cache.get_cache(key=cache_key) - # wait 2 seconds time.sleep(2) + return_dict = { + deployment_id: router.cache.get_cache(key=cache_key) + for deployment_id, cache_key in cache_keys.items() + } + assert router.leastbusy_logger.logged_success == 1 - assert return_dict[1] == 10 - assert return_dict[2] == 54 - assert return_dict[3] == 100 + assert return_dict["1"] == 10 + assert return_dict["2"] == 54 + assert return_dict["3"] == 100 ## Test with Real calls ## @@ -192,9 +200,11 @@ async def test_router_atext_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.atext_completion(model=model, prompt=prompt, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" @@ -259,8 +269,10 @@ async def test_router_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.acompletion(model=model, messages=messages, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" diff --git a/tests/local_testing/test_lowest_latency_routing.py b/tests/local_testing/test_lowest_latency_routing.py index 598b1dbcaf9..aba4500199a 100644 --- a/tests/local_testing/test_lowest_latency_routing.py +++ b/tests/local_testing/test_lowest_latency_routing.py @@ -1077,73 +1077,6 @@ async def test_latency_list_trimming_discards_oldest_entry_async(): ), f"Oldest latency {oldest_latency} should have been discarded" -def test_ttft_list_trimming_discards_oldest_entry(): - """ - The time_to_first_token list trims the oldest entry when full, matching - the behavior of the latency list. - """ - max_size = 3 - test_cache = DualCache() - lowest_latency_logger = LowestLatencyLoggingHandler( - router_cache=test_cache, routing_args={"max_latency_list_size": max_size} - ) - - model_group = "gpt-3.5-turbo" - deployment_id = "test-deployment" - - ttft_values = [] - for i in range(max_size + 1): - start_time = time.time() - expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4 - completion_start_time = start_time + expected_ttft - end_time = start_time + float(i + 1) - ttft_values.append(expected_ttft) - - kwargs = { - "litellm_params": { - "metadata": { - "model_group": model_group, - "deployment": "azure/gpt-4.1-mini", - }, - "model_info": {"id": deployment_id}, - }, - "stream": True, - "completion_start_time": completion_start_time, - } - # TTFT is only recorded when response_obj is a ModelResponse. - response_obj = litellm.ModelResponse( - usage=litellm.Usage(completion_tokens=1, total_tokens=1) - ) - - lowest_latency_logger.log_success_event( - response_obj=response_obj, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - - latency_key = f"{model_group}_map" - cached_data = test_cache.get_cache(key=latency_key) - ttft_list = cached_data[deployment_id].get("time_to_first_token", []) - - assert ( - len(ttft_list) == max_size - ), f"Expected {max_size} entries, got {len(ttft_list)}" - - newest_ttft = ttft_values[-1] - oldest_ttft = ttft_values[0] - tolerance = 0.05 - - assert ( - abs(ttft_list[-1] - newest_ttft) < tolerance - ), f"Newest TTFT {newest_ttft} should be at end of list" - - for ttft in ttft_list: - assert ( - abs(ttft - oldest_ttft) > tolerance - ), f"Oldest TTFT {oldest_ttft} should have been discarded" - - @pytest.mark.asyncio async def test_timeout_penalty_discards_oldest_entry(): """ @@ -1269,72 +1202,3 @@ def test_list_order_preserved_after_multiple_trims(): assert ( abs(latency_list[i] - expected) < tolerance ), f"At index {i}, expected ~{expected}, got {latency_list[i]}" - - -@pytest.mark.asyncio -async def test_ttft_list_trimming_discards_oldest_entry_async(): - """ - Async counterpart: the time_to_first_token list trims the oldest entry - when full. Exercises the async_log_success_event TTFT path, which only - runs when response_obj is a ModelResponse and the call is marked as - streaming with a completion_start_time. - """ - max_size = 3 - test_cache = DualCache() - lowest_latency_logger = LowestLatencyLoggingHandler( - router_cache=test_cache, routing_args={"max_latency_list_size": max_size} - ) - - model_group = "gpt-3.5-turbo" - deployment_id = "test-deployment" - - ttft_values = [] - for i in range(max_size + 1): - start_time = time.time() - expected_ttft = float(i + 1) * 0.1 # 0.1, 0.2, 0.3, 0.4 - completion_start_time = start_time + expected_ttft - end_time = start_time + float(i + 1) - ttft_values.append(expected_ttft) - - kwargs = { - "litellm_params": { - "metadata": { - "model_group": model_group, - "deployment": "azure/gpt-4.1-mini", - }, - "model_info": {"id": deployment_id}, - }, - "stream": True, - "completion_start_time": completion_start_time, - } - response_obj = litellm.ModelResponse( - usage=litellm.Usage(completion_tokens=1, total_tokens=1) - ) - - await lowest_latency_logger.async_log_success_event( - response_obj=response_obj, - kwargs=kwargs, - start_time=start_time, - end_time=end_time, - ) - - latency_key = f"{model_group}_map" - cached_data = await test_cache.async_get_cache(key=latency_key) - ttft_list = cached_data[deployment_id].get("time_to_first_token", []) - - assert ( - len(ttft_list) == max_size - ), f"Expected {max_size} entries, got {len(ttft_list)}" - - newest_ttft = ttft_values[-1] - oldest_ttft = ttft_values[0] - tolerance = 0.05 - - assert ( - abs(ttft_list[-1] - newest_ttft) < tolerance - ), f"Newest TTFT {newest_ttft} should be at end of list" - - for ttft in ttft_list: - assert ( - abs(ttft - oldest_ttft) > tolerance - ), f"Oldest TTFT {oldest_ttft} should have been discarded" diff --git a/tests/local_testing/test_redis_increment_with_floor.py b/tests/local_testing/test_redis_increment_with_floor.py new file mode 100644 index 00000000000..e358d5f31e0 --- /dev/null +++ b/tests/local_testing/test_redis_increment_with_floor.py @@ -0,0 +1,80 @@ +"""Least-busy routing keeps its in-flight counters in Redis, and the clamp at zero plus the +create-once TTL both live inside a Lua script. Nothing but a real Redis runs that script, so +these are the only tests that fail when the script itself is wrong.""" + +import os +import uuid +from typing import Final + +import pytest +from dotenv import load_dotenv + +load_dotenv() + +from litellm.caching.redis_cache import RedisCache + +TTL: Final = 600 + + +@pytest.fixture +def counter(): + cache: Final = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT")) + key: Final = f"lit7039-{uuid.uuid4()}" + yield cache, key, cache.check_and_fix_namespace(key=key) + cache.delete_cache(key) + + +def test_a_counter_adds_every_increment_and_reads_back_what_it_holds(counter): + cache, key, _ = counter + + assert cache.increment_with_floor(key, 3, TTL) == 3 + assert cache.increment_with_floor(key, 2, TTL) == 5 + assert cache.batch_get_counts([key]) == (5,) + + +def test_a_decrement_past_zero_leaves_the_counter_at_zero(counter): + """A worker whose counter expired mid-request decrements a key that is no longer there. + Without the clamp that deployment reads negative, and least-busy pins every later request + on it until the count climbs back to zero.""" + cache, key, _ = counter + + assert cache.increment_with_floor(key, 1, TTL) == 1 + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.batch_get_counts([key]) == (0,) + + +def test_traffic_never_pushes_a_counters_expiry_back_out(counter): + """The TTL is what releases a count whose worker died mid-request. Rewriting it on every + touch would keep that stuck count alive for as long as the group takes traffic.""" + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + assert cache.redis_client.ttl(namespaced_key) > TTL - 60 + + cache.redis_client.expire(namespaced_key, 30) + cache.increment_with_floor(key, 1, TTL) + + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter): + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + cache.redis_client.expire(namespaced_key, 30) + + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +@pytest.mark.asyncio +async def test_the_async_counter_behaves_the_same_way(counter): + cache, key, namespaced_key = counter + + assert await cache.async_increment_with_floor(key, 2, TTL) == 2 + assert await cache.async_batch_get_counts([key]) == (2,) + + cache.redis_client.expire(namespaced_key, 30) + + assert await cache.async_increment_with_floor(key, -9, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index 7bc61c40ea0..9ac6476a03c 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -40,12 +40,26 @@ async def _turn( tokens: int = 100, spend: float = 0.01, saved: float = 0.02, + classifier_cost: float = 0.0, tier: "str | None" = None, ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( UPSERT_AUTOROUTER_SESSION_SQL, - key, session_id, router, router_type, model, at.isoformat(), tokens, spend, saved, covered, hit, ttl, touched, + key, + session_id, + router, + router_type, + model, + at.isoformat(), + tokens, + spend, + saved, + classifier_cost, + covered, + hit, + ttl, + touched, tier, ) @@ -53,7 +67,9 @@ async def _turn( async def _row(db, key: str, session_id: str = "s1", router: str = "auto-1") -> dict: rows = await db.query_raw( 'SELECT * FROM "LiteLLM_AutoRouterSession" WHERE api_key = $1 AND session_id = $2 AND router_name = $3', - key, session_id, router, + key, + session_id, + router, ) assert len(rows) == 1 return rows[0] @@ -143,9 +159,9 @@ async def test_out_of_order_turns_do_not_rewind_the_session(db): async def test_concurrent_writers_compose_without_losing_turns(db): key = f"k-{uuid.uuid4()}" - await _turn(db, key, "A", T0) + await _turn(db, key, "A", T0, classifier_cost=0.001) await asyncio.gather( - *(_turn(db, key, "A", T0 + timedelta(seconds=1 + offset), hit=1) for offset in range(30)) + *(_turn(db, key, "A", T0 + timedelta(seconds=1 + offset), hit=1, classifier_cost=0.002) for offset in range(30)) ) row = await _row(db, key) assert row["turns"] == 31 @@ -154,6 +170,51 @@ async def test_concurrent_writers_compose_without_losing_turns(db): == row["turns"] ) assert row["spend"] == pytest.approx(0.31) + assert row["saved_spend"] == pytest.approx(0.62) + assert row["classifier_cost"] == pytest.approx(0.061) + assert row["classifier_cost_recorded_turns"] == 31 + + +async def _legacy_turn(db, key: str, at: datetime, session_id: str = "s1", router: str = "auto-1") -> None: + await db.execute_raw( + """INSERT INTO "LiteLLM_AutoRouterSession" AS t ( + api_key, session_id, router_name, router_type, first_turn_at, last_turn_at, last_model, turns, spend, saved_spend + ) VALUES ($1, $2, $3, 'complexity', $4::timestamp, $4::timestamp, 'A', 1, 0.01, 0.02) + ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET + turns = t.turns + 1, spend = t.spend + EXCLUDED.spend, saved_spend = t.saved_spend + EXCLUDED.saved_spend, + last_turn_at = EXCLUDED.last_turn_at""", + key, + session_id, + router, + at.isoformat(), + ) + + +@pytest.mark.parametrize("writers", [(False,), (True,), (False, True), (True, False)]) +async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers: tuple[bool, ...]): + key: Final = f"k-{uuid.uuid4()}" + for offset, records_cost in enumerate(writers): + at: Final = T0 + timedelta(seconds=offset) + if records_cost: + await _turn(db, key, "A", at, classifier_cost=0.004) + else: + await _legacy_turn(db, key, at) + + row: Final = await _row(db, key) + assert row["turns"] == len(writers) + assert row["spend"] == pytest.approx(0.01 * len(writers)) + assert row["saved_spend"] == pytest.approx(0.02 * len(writers)) + assert row["classifier_cost"] == pytest.approx(0.004 * sum(writers)) + assert row["classifier_cost_recorded_turns"] == sum(writers) + groups: Final = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key + ) + assert len(groups) == 1 + assert groups[0]["classifier_cost"] == row["classifier_cost"] + assert groups[0]["classifier_cost_recorded_turns"] == sum(writers) + assert groups[0]["turns"] == len(writers) + assert groups[0]["spend"] == row["spend"] + assert groups[0]["saved_spend"] == row["saved_spend"] async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): @@ -161,14 +222,25 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): router = f"r-{uuid.uuid4()}" in_window = f"s-{uuid.uuid4()}" out_of_window = f"s-{uuid.uuid4()}" - await _turn(db, key, "A", T0, session_id=in_window, router=router, saved=0.5, spend=0.25) - await _turn(db, key, "B", T0 + timedelta(seconds=60), session_id=in_window, router=router, saved=0.5, spend=0.25) - await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router) + await _turn( + db, + key, + "B", + T0 + timedelta(seconds=60), + session_id=in_window, + router=router, + saved=0.5, + spend=0.25, + classifier_cost=0.01, + ) + await _turn(db, key, "A", T0, session_id=in_window, router=router, saved=0.5, spend=0.25, classifier_cost=0.02) + await _turn(db, key, "A", T0 - timedelta(days=40), session_id=out_of_window, router=router, classifier_cost=9.0) rows = await db.query_raw( 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 @@ -178,19 +250,54 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): assert grouped["turns"] == 2 assert grouped["spend"] == pytest.approx(0.5) assert grouped["saved_spend"] == pytest.approx(1.0) + assert grouped["classifier_cost"] == pytest.approx(0.03) + assert grouped["classifier_cost_recorded_turns"] == 2 + assert grouped["unordered_turns"] == 1 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, router=router, saved=0.5, classifier_cost=0.01) + await _turn(db, second_key, "A", T0, router=router, saved=9.0, classifier_cost=0.09) + + 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) + assert matching[0]["classifier_cost"] == pytest.approx(0.01) + assert matching[0]["classifier_cost_recorded_turns"] == 1 + + 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()}" await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity") - await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality") + await _turn( + db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, router_type="quality" + ) rows = await db.query_raw( 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 +360,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 +388,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 +403,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/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py new file mode 100644 index 00000000000..741fa7386df --- /dev/null +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -0,0 +1,264 @@ +import os +import threading +import uuid +from collections.abc import Iterator, Mapping +from types import MappingProxyType +from typing import Final + +import pytest +from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager + +psycopg = pytest.importorskip("psycopg") + +pytestmark = pytest.mark.timeout(120) + +requires_db: Final = pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) + +HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable" +HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx" +HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC' +LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable" +LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx" +PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable" +PARTITIONED_INDEX: Final = "LiteLLM_PartitionedTable_id_idx" + + +def _base_url() -> str: + return os.environ["DATABASE_URL"].split("?")[0] + + +def _index_validity(schema: str) -> Mapping[str, bool]: + with psycopg.connect(_base_url(), autocommit=True) as conn: + rows = conn.execute( + "SELECT c.relname, i.indisvalid FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_namespace n ON n.oid = c.relnamespace " + "WHERE n.nspname = %s", + (schema,), + ).fetchall() + return MappingProxyType(dict(rows)) + + +def _interrupt_concurrent_build(schema: str, table: str, statement: str) -> None: + """Abort a CONCURRENTLY build while it waits on an older snapshot, the same + spot the deadlock loser dies at, so it leaves its index INVALID.""" + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"') + with psycopg.connect(_base_url(), autocommit=True) as builder: + builder.execute("SET statement_timeout = '1s'") + with pytest.raises(psycopg.errors.QueryCanceled): + builder.execute(statement) + + +def _leave_invalid_index(schema: str, table: str, index: str, columns: str) -> None: + _interrupt_concurrent_build( + schema, table, f'CREATE INDEX CONCURRENTLY "{index}" ON "{schema}"."{table}" ({columns})' + ) + + +def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None: + _interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"') + + +@pytest.fixture +def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE SCHEMA "{schema}"') + conn.execute( + f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)' + ) + conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)') + + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") + yield schema + + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP SCHEMA "{schema}" CASCADE') + + +@pytest.fixture +def fresh_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + """A brand-new database, what a first deploy sees. A scratch schema would + not do: the migrations guard on pg_constraint by name across every schema, + so a LiteLLM schema already pushed into public makes them skip and then + fail, which is exactly what CI's database looks like.""" + admin_url: Final = _base_url() + name: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'CREATE DATABASE "{name}"') + + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{admin_url.rsplit('/', 1)[0]}/{name}") + yield "public" + + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)') + + +@requires_db +def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + _leave_invalid_index(scratch_schema, LOOKALIKE_TABLE, LOOKALIKE_INDEX, "id") + assert _index_validity(scratch_schema) == {HEALTH_INDEX: False, LOOKALIKE_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, LOOKALIKE_INDEX: False} + + +@requires_db +def test_repair_drops_leftovers_of_interrupted_rebuilds(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + _leave_invalid_reindex_leftover(scratch_schema, HEALTH_TABLE, HEALTH_INDEX) + _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccold", '"model_id"') + _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccnew1", '"model_id"') + before: Final = _index_validity(scratch_schema) + assert len(before) == 4 + assert set(before.values()) == {False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_is_a_no_op_when_every_index_is_valid(scratch_schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE INDEX "{HEALTH_INDEX}" ON "{scratch_schema}"."{HEALTH_TABLE}" ({HEALTH_INDEX_COLUMNS})') + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_leaves_partitioned_parent_indexes_alone(scratch_schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}" (id INT) PARTITION BY RANGE (id)') + conn.execute( + f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}_p0" ' + f'PARTITION OF "{scratch_schema}"."{PARTITIONED_TABLE}" FOR VALUES FROM (0) TO (10)' + ) + conn.execute(f'CREATE INDEX "{PARTITIONED_INDEX}" ON ONLY "{scratch_schema}"."{PARTITIONED_TABLE}" (id)') + assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False} + + +@requires_db +def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + with psycopg.connect(_base_url(), autocommit=True) as other_replica: + other_replica.execute("SELECT pg_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)) + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + assert _index_validity(scratch_schema) == {HEALTH_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{scratch_schema}"."{HEALTH_TABLE}"') + assert ProxyExtrasDBManager.repair_invalid_indexes(lock_timeout="1s") is False + blocked: Final = _index_validity(scratch_schema) + assert blocked[HEALTH_INDEX] is False + assert [name for name in blocked if name.endswith("_ccnew")] + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +def _hold_snapshot(schema: str, table: str, pinned: threading.Event, seconds: float) -> None: + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"') + pinned.set() + pin.execute("SELECT pg_sleep(%s)", (seconds,)) + + +@requires_db +def test_repair_outlives_a_statement_timeout_passed_through_database_url_options( + scratch_schema: str, monkeypatch: pytest.MonkeyPatch +) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={scratch_schema}&options=-c%20statement_timeout%3D2000") + pinned: Final = threading.Event() + holder: Final = threading.Thread(target=_hold_snapshot, args=(scratch_schema, HEALTH_TABLE, pinned, 5.0)) + holder.start() + pinned.wait() + try: + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + finally: + holder.join() + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) -> None: + table: Final = f"LiteLLM_ScratchTable_{uuid.uuid4().hex[:8]}" + index: Final = f"{table}_id_idx" + monkeypatch.setenv("DATABASE_URL", _base_url()) + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE TABLE public."{table}" (id TEXT)') + try: + _leave_invalid_index("public", table, index, "id") + assert _index_validity("public")[index] is False + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity("public")[index] is True + finally: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP TABLE public."{table}"') + + +def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + + +@requires_db +def test_repair_connects_over_direct_url_but_looks_in_the_schema_database_url_names(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + with pytest.MonkeyPatch.context() as env: + env.setenv("DIRECT_URL", f"{_base_url()}?schema=public") + env.setenv("DATABASE_URL", f"postgresql://u:p@127.0.0.1:9/x?schema={scratch_schema}") + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +def _invalidate_deployed_index(schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP INDEX "{schema}"."{HEALTH_INDEX}"') + _leave_invalid_index(schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + +@requires_db +@pytest.mark.timeout(300) +@pytest.mark.parametrize("use_v2_resolver", [True, False]) +def test_setup_database_repairs_the_index_after_a_recovered_deploy(fresh_database: str, use_v2_resolver: bool) -> None: + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + _invalidate_deployed_index(fresh_database) + assert _index_validity(fresh_database)[HEALTH_INDEX] is False + + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + + assert _index_validity(fresh_database)[HEALTH_INDEX] is True diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 9a6ab08e9b6..dddc8304bfc 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1760,6 +1760,35 @@ class TestUnmanagedVertexRouting: ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") + def test_flag_on_routes_fine_tuned_endpoint_to_vertex_deployment(self): + """A fine-tuned Gemini batch stores `endpoints/` in the gs:// path; the bare model + (the endpoint id) must round-trip to the deployment configured as + `vertex_ai/gemini/` (LIT-6899).""" + endpoint_id = "7768560373388541952" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_list.return_value = [ + { + "model_name": "gemini-2.5-flash-dts-usc1", + "litellm_params": { + "model": f"vertex_ai/gemini/{endpoint_id}", + "custom_llm_provider": "vertex_ai", + }, + "model_info": {"id": "deploy-ft"}, + }, + ] + instance = self._instance(track_unmanaged=True, router=router) + job = self._job( + file_object=_unmanaged_vertex_file_object( + input_file_id=f"gs://bucket/litellm-vertex-files/endpoints/{endpoint_id}/abc.jsonl" + ) + ) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-ft", "8823717160934178816") + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): """Flag on, but the only deployment for the model group is a non-vertex_ai provider: must not be selected, even though the model group name matches.""" diff --git a/tests/router_unit_tests/test_router_cooldown_per_deployment.py b/tests/router_unit_tests/test_router_cooldown_per_deployment.py index b8ae8a8c013..964782c348b 100644 --- a/tests/router_unit_tests/test_router_cooldown_per_deployment.py +++ b/tests/router_unit_tests/test_router_cooldown_per_deployment.py @@ -204,8 +204,8 @@ class TestExceptionTypeCountersTrackedIndependently: cache_key_suffix="RateLimitError", ) - rl_counter = router.failed_calls.get_cache(key="primary:RateLimitError") or 0 - generic_counter = router.failed_calls.get_cache(key="primary:generic") or 0 + rl_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0 + generic_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0 assert rl_counter == 3, "RateLimitError counter should be 3" assert generic_counter == 0, "generic counter must be untouched by RateLimitError increments" @@ -218,8 +218,8 @@ class TestExceptionTypeCountersTrackedIndependently: cache_key_suffix="generic", ) - generic_counter_after = router.failed_calls.get_cache(key="primary:generic") or 0 - rl_counter_after = router.failed_calls.get_cache(key="primary:RateLimitError") or 0 + generic_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0 + rl_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0 assert generic_counter_after == 1, "generic counter should now be 1" assert rl_counter_after == 3, "RateLimitError counter must remain unchanged after InternalServerError" @@ -246,12 +246,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired cooldown entry must not appear in active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" + assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" def test_active_entry_is_returned(self): """ @@ -267,7 +267,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -290,14 +290,14 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - (60.0 - remaining), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + cc.in_memory_cache.set_cache(key, value, ttl=600) - before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + before_expiry = cc.in_memory_cache.ttl_dict.get(key) assert before_expiry is not None cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry is not None corrected_remaining = after_expiry - time.time() assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" @@ -318,12 +318,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired entry must not appear in async active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None class TestFallbackDeploymentCooldown: diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index c3edb40c819..b87f9489250 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -158,6 +158,14 @@ def test_create__vertex_ai_dispatch(seams): _assert_only(seams.vertex.create_batch, seams, "create_batch") +def test_create__vertex_ai_forwards_custom_endpoint(seams): + """The vertex handler owns the custom_endpoint batch rejection (LIT-6899), so the dispatcher + must forward the flag for the handler to act on.""" + bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True) + + assert seams.vertex.create_batch.call_args.kwargs["custom_endpoint"] is True + + def test_create__provider_config_routes_to_base_http_handler(seams): """model + a provider batches config (bedrock-style) routes to the generic base_llm_http_handler, NOT the per-provider instance.""" diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 2f412e7382b..6b2df118611 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -525,6 +525,50 @@ def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_re assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} +def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache): + """A caller that must fall back when Redis is unreachable needs the failure, not zeros. + + The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell + apart from "every counter is unset". Least-busy routing read that as an idle deployment and + kept sending traffic to it instead of falling back to this worker's own in-flight counts. + """ + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + sync_batch_redis_cache.batch_get_counts(["lit7039"]) + + +@pytest.mark.asyncio +async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None): + """Async twin: the async batch read hides the same failure behind an empty dict.""" + failing_client = AsyncMock() + failing_client.mget.side_effect = OSError("redis unavailable") + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + + with patch.object(cache, "init_async_client", return_value=failing_client): + assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + await cache.async_batch_get_counts(["lit7039"]) + + +@pytest.mark.parametrize("stored", [b"3", "3"]) +def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None): + """Counters come back positionally, so an unset key has to stay a hole rather than shift the + rest of the row onto the wrong deployments, and a count has to survive whether the client + hands it back as bytes or as text.""" + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client.mget.return_value = [stored, None, b"0"] + + assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) + + @pytest.fixture def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]: service_logger = ServiceLogging(mock_testing=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 d4d47b145d1..e4ada0a9b31 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 @@ -4147,3 +4147,143 @@ def test_streaming_final_chunk_carries_provider_metadata(): 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]) + + +def _system_input_item(text: str) -> dict[str, object]: + return {"type": "message", "role": "system", "content": [{"type": "input_text", "text": text}]} + + +def test_mid_conversation_system_string_stays_in_input_after_a_user_turn(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Read the file."}, + {"role": "system", "content": "14982391 tokens left"}, + {"role": "user", "content": "Now summarize it."}, + ] + ) + + assert instructions == "You are a helpful assistant." + assert input_items == [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Read the file."}]}, + _system_input_item("14982391 tokens left"), + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Now summarize it."}]}, + ] + + +def test_leading_system_strings_still_join_instructions_without_a_following_turn(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "system", "content": "Be brief."}, + {"role": "system", "content": "Answer in French."}, + ] + ) + + assert instructions == "Be brief. Answer in French." + assert input_items == [] + + +def test_mid_conversation_system_reminder_as_string_and_as_text_block_produce_identical_input_items(): + handler: Final = LiteLLMResponsesTransformationHandler() + reminder: Final = "14982391 tokens left" + + as_string, string_instructions = handler.convert_chat_completion_messages_to_responses_api( + [{"role": "user", "content": "Read the file."}, {"role": "system", "content": reminder}] + ) + as_block, block_instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "user", "content": "Read the file."}, + { + "role": "system", + "content": [{"type": "text", "text": reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + ) + + assert string_instructions is None + assert block_instructions is None + assert json.dumps(as_string) == json.dumps(as_block) + assert as_string[1] == _system_input_item(reminder) + + +def test_claude_code_shaped_history_keeps_a_byte_stable_input_prefix_across_requests(): + handler: Final = LiteLLMResponsesTransformationHandler() + top_level_system: Final = [{"type": "text", "text": "You are Claude Code.", "cache_control": {"type": "ephemeral"}}] + first_reminder: Final = "27k chars of deferred tools" + second_reminder: Final = "14982391 tokens left" + first_request_messages: Final = [ + {"role": "system", "content": top_level_system}, + {"role": "user", "content": "Read inventory.py."}, + { + "role": "system", + "content": [{"type": "text", "text": first_reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + second_request_messages: Final = [ + {"role": "system", "content": top_level_system}, + {"role": "user", "content": "Read inventory.py."}, + {"role": "system", "content": first_reminder}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "Read", "arguments": '{"file_path": "inventory.py"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "ITEMS = []"}, + { + "role": "system", + "content": [{"type": "text", "text": second_reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + + first_request: Final = handler.transform_request( + model="gpt-5.6-luna", + messages=first_request_messages, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + second_request: Final = handler.transform_request( + model="gpt-5.6-luna", + messages=second_request_messages, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert "instructions" not in first_request + assert "instructions" not in second_request + assert first_request["input"][0] == _system_input_item("You are Claude Code.") + assert json.dumps(second_request["input"][: len(first_request["input"])]) == json.dumps(first_request["input"]) + assert second_request["input"][len(first_request["input"]) :] == [ + {"type": "function_call", "call_id": "call_1", "name": "Read", "arguments": '{"file_path": "inventory.py"}'}, + {"type": "function_call_output", "call_id": "call_1", "output": [{"type": "input_text", "text": "ITEMS = []"}]}, + _system_input_item(second_reminder), + ] + + +def test_system_string_after_a_developer_message_stays_in_input_in_client_order(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "developer", "content": "Always answer in French."}, + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Bonjour"}, + ] + ) + + assert instructions is None + assert [item["role"] for item in input_items] == ["developer", "system", "user"] + assert input_items[1] == _system_input_item("Be brief.") diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py index f46df5baadf..c1c3569c0e2 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -262,7 +262,7 @@ def test_proxied_traffic_stays_on_native_hooks(): never sees ``data["prompt"]``.""" guardrail = _guardrail() assert guardrail.uses_apply_guardrail_interface() is True - assert guardrail._deployment_pre_call_target() is guardrail + assert guardrail._deployment_hook_target() is guardrail @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 4aa28b5abfd..6676714d7f5 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -3,7 +3,10 @@ baggage helpers, metrics, the typed coercion helpers, mapper branches, span-name builders, and the registry validator's failure paths. Needs the OTel SDK.""" import json +import threading +from collections.abc import Iterator from dataclasses import replace +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest @@ -538,6 +541,77 @@ def test_build_span_exporter_variants(): assert "OTLPSpanExporter" in type(http_exporter).__name__ +@pytest.fixture +def otlp_collector() -> Iterator[tuple[str, list[str]]]: + received_paths: list[str] = [] + + class RecordingHandler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + self.rfile.read(int(self.headers.get("Content-Length", "0"))) + received_paths.append(self.path) + self.send_response(200) + self.end_headers() + + def log_message(self, format: str, *args: object) -> None: + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), RecordingHandler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + yield f"http://127.0.0.1:{server.server_port}", received_paths + finally: + server.shutdown() + server.server_close() + + +def _export_one_span(cfg: OpenTelemetryV2Config) -> None: + provider = providers.build_tracer_provider(cfg) + provider.get_tracer("probe").start_span("probe").end() + assert provider.force_flush() + provider.shutdown() + + +def test_traces_endpoint_env_posts_to_the_configured_url_verbatim(monkeypatch, otlp_collector): + base_url, received_paths = otlp_collector + for var in ("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_EXPORTER_OTLP_ENDPOINT"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("OTEL_ENDPOINT", f"{base_url}/services/collector") + monkeypatch.setenv("OTEL_TRACES_ENDPOINT", f"{base_url}/services/collector/traces") + + cfg = OpenTelemetryV2Config.from_env() + assert cfg.exporter == "otlp_http" + _export_one_span(cfg) + assert received_paths == ["/services/collector/traces"] + + +def test_traces_endpoint_alias_alone_implies_otlp_http(monkeypatch, otlp_collector): + base_url, received_paths = otlp_collector + for var in ("OTEL_EXPORTER", "OTEL_EXPORTER_OTLP_PROTOCOL", "OTEL_ENDPOINT", "OTEL_EXPORTER_OTLP_ENDPOINT"): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", f"{base_url}/custom/traces") + + cfg = OpenTelemetryV2Config.from_env() + assert cfg.exporter == "otlp_http" + _export_one_span(cfg) + assert received_paths == ["/custom/traces"] + + +def test_traces_endpoint_per_exporter_coexists_with_default_normalization(otlp_collector): + base_url, received_paths = otlp_collector + cfg = OpenTelemetryV2Config( + exporters=[ + {"kind": "otlp_http", "endpoint": base_url}, + { + "kind": "otlp_http", + "endpoint": f"{base_url}/services/collector", + "traces_endpoint": f"{base_url}/services/collector/traces", + }, + ] + ) + _export_one_span(cfg) + assert sorted(received_paths) == ["/services/collector/traces", "/v1/traces"] + + def test_otlp_metric_exporter_uses_cumulative_histogram_temporality(): """Histograms must export as cumulative, not delta. diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index de8b654987b..6b7780acd20 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1777,6 +1777,224 @@ class TestEnableAnthropicPromptCaching: assert messages == before +class TestClaudeCodeOneShotAutoCaching: + BILLING_TEXT = "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli; cc_is_subagent=true;" + BILLING_SYSTEM = [{"type": "text", "text": BILLING_TEXT}] + MESSAGES = [{"role": "user", "content": [{"type": "text", "text": "unique fetched document"}]}] + + @staticmethod + def _kwargs(configured=None): + kwargs = { + "litellm_metadata": {}, + "proxy_server_request": { + "headers": { + "user-agent": "claude-cli/2.1.263 (external, cli)", + "x-app": "cli-bg", + } + }, + } + if configured is not None: + kwargs["cache_control_injection_points"] = configured + return kwargs + + @pytest.mark.parametrize( + "system", + [ + BILLING_TEXT, + BILLING_SYSTEM, + [*BILLING_SYSTEM, {"type": "text", "text": " "}], + [ + *BILLING_SYSTEM, + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_entrypoint=cli;"}, + ], + ], + ids=["string", "text_block", "whitespace_block", "multiple_billing_blocks"], + ) + @pytest.mark.parametrize("tools", [None, []], ids=["absent_tools", "empty_tools"]) + def test_skips_defaults_and_attribution_for_one_shot_subagent(self, monkeypatch, system, tools): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = copy.deepcopy(self.MESSAGES) + kwargs = self._kwargs() + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + messages, + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=tools, + ) + + assert result_messages == self.MESSAGES + assert result_system == system + assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"] + + def test_user_agent_header_lookup_is_case_insensitive(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + user_agent = kwargs["proxy_server_request"]["headers"].pop("user-agent") + kwargs["proxy_server_request"]["headers"]["User-Agent"] = user_agent + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(self.BILLING_SYSTEM), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert result_messages == self.MESSAGES + assert result_system == self.BILLING_SYSTEM + + def test_router_affinity_skips_string_billing_system(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages = copy.deepcopy(self.MESSAGES) + kwargs = self._kwargs() + kwargs["system"] = self.BILLING_TEXT + + result = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=("claude-sonnet-4-5",), + request_kwargs=kwargs, + ) + + assert result == messages + + @pytest.mark.parametrize( + "headers,system", + [ + ("not-a-mapping", BILLING_SYSTEM), + ( + {"user-agent": "claude-cli/2.1.263 (external, cli)"}, + [{"type": "text", "text": "x-anthropic-billing-header: malformed"}], + ), + ({"user-agent": "claude-cli/2.1.263 (external, cli)"}, None), + ({"user-agent": "claude-cli/2.1.263 (external, cli)"}, ["not-a-mapping"]), + ( + {"user-agent": "claude-cli/2.1.263 (external, cli)"}, + [{"type": "image", "text": BILLING_TEXT}], + ), + ], + ids=["malformed_headers", "malformed_billing", "missing_system", "malformed_block", "non_text_block"], + ) + def test_malformed_untrusted_context_keeps_defaults(self, monkeypatch, headers, system): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + points = AnthropicCacheControlHook.get_default_injection_points( + messages=copy.deepcopy(self.MESSAGES), + system=copy.deepcopy(system), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + request_kwargs={"proxy_server_request": {"headers": headers}}, + ) + + assert len(points) == 2 + + def test_message_without_role_keeps_defaults(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + points = AnthropicCacheControlHook.get_default_injection_points( + messages=[{"content": "missing role"}], + system=copy.deepcopy(self.BILLING_SYSTEM), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + request_kwargs=self._kwargs(), + ) + + assert len(points) == 2 + + @pytest.mark.parametrize( + "messages,system,tools", + [ + ( + MESSAGES, + BILLING_SYSTEM, + [{"name": "WebFetch", "description": "fetch", "input_schema": {"type": "object"}}], + ), + (MESSAGES, [*BILLING_SYSTEM, {"type": "text", "text": "Explore the repository"}], None), + ( + [ + {"role": "user", "content": "first turn"}, + {"role": "assistant", "content": "reply"}, + *MESSAGES, + ], + BILLING_SYSTEM, + None, + ), + ], + ids=["tools", "real_system", "history"], + ) + def test_keeps_defaults_for_reusable_subagents(self, monkeypatch, messages, system, tools): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(messages), + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + tools=copy.deepcopy(tools), + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2 + assert kwargs["litellm_metadata"]["litellm_gateway_injected_cache"] == "" + + @pytest.mark.parametrize( + "user_agent,system", + [ + ("anthropic-sdk-python/0.75.0", BILLING_SYSTEM), + ( + "claude-cli/2.1.263 (external, cli)", + [ + { + "type": "text", + "text": f"{BILLING_TEXT}\nadditional system instructions", + } + ], + ), + ( + "claude-cli/2.1.263 (external, cli)", + [ + { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=false;", + } + ], + ), + ], + ids=["different_client", "appended_instructions", "not_a_subagent"], + ) + def test_ambiguous_or_unmatched_signals_fail_open(self, monkeypatch, user_agent, system): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs() + kwargs["proxy_server_request"]["headers"]["user-agent"] = user_agent + + result_messages, result_system = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(system), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert AnthropicCacheControlHook.count_request_cache_breakpoints(result_messages, result_system) == 2 + + def test_explicit_injection_points_remain_authoritative(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + kwargs = self._kwargs([{"location": "message", "role": "user"}]) + + result_messages, _ = AnthropicCacheControlHook.maybe_inject_cache_control( + copy.deepcopy(self.MESSAGES), + copy.deepcopy(self.BILLING_SYSTEM), + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + + assert result_messages[0]["content"][-1]["cache_control"] == {"type": "ephemeral"} + + class TestPerKeyEnablePromptCaching: """Per-request enable_prompt_caching override (stamped from key metadata) with the global flag off.""" @@ -1977,6 +2195,25 @@ class TestConfiguredInjectionPointsStandDown: _, result_sys = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) assert result_sys == [{"type": "text", "text": "sys", "cache_control": {"type": "ephemeral"}}] + @pytest.mark.parametrize( + "configured", + [None, CONFIGURED], + ids=["automatic_defaults", "configured_points"], + ) + def test_v1_messages_stands_down_for_root_cache_control(self, monkeypatch, configured): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + root_cache_control = {"type": "ephemeral"} + kwargs = {"cache_control": root_cache_control, "litellm_metadata": {}} + if configured is not None: + kwargs["cache_control_injection_points"] = copy.deepcopy(configured) + + result_messages, result_system = self._inject(copy.deepcopy(self.V1_MESSAGES), kwargs) + + assert result_messages == self.V1_MESSAGES + assert result_system == "sys" + assert kwargs["cache_control"] is root_cache_control + assert "litellm_gateway_injected_cache" not in kwargs["litellm_metadata"] + def test_v1_messages_reentry_flow_preserves_tool_config_remainder(self): """The advisor interceptor re-enters anthropic_messages() with the outer request's kwargs and post-injection messages. The first pass applies the diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 885bd1d4d72..cd8d609cf71 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Literal, Optional +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock import pytest @@ -10,6 +10,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import CallTypes, UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail if TYPE_CHECKING: @@ -2378,11 +2379,108 @@ def _logged_call(messages: list | str) -> tuple[dict, object]: return kwargs, response +class _NativeApplyGuardrail(_InheritedApplyGuardrail): + use_native_lifecycle_hooks: ClassVar[bool] = True + + +@pytest.mark.parametrize("guardrail_type", (CustomGuardrail, _NativeApplyGuardrail, _InheritedApplyGuardrail)) +@pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.logging_only, + "logging_only", + [GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ["pre_call", "logging_only"], + Mode(tags={"audit": "logging_only"}, default="pre_call"), + Mode(tags={"audit": ["pre_call", "logging_only"]}), + Mode(tags={"enforce": "pre_call"}, default="logging_only"), + Mode(tags={}, default=["pre_call", "logging_only"]), + ), +) +def test_logging_only_requires_framework_support_or_explicit_declaration( + guardrail_type: type[CustomGuardrail], + event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, +) -> None: + supported: Final = [GuardrailEventHooks.pre_call] + if guardrail_type is _InheritedApplyGuardrail: + guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) + assert guardrail.event_hook == event_hook + assert supported == [GuardrailEventHooks.pre_call] + else: + with pytest.raises(ValueError, match=r"logging_only.*not in the supported event hooks"): + guardrail_type(event_hook=event_hook, supported_event_hooks=supported) + + explicitly_supported: Final = guardrail_type( + event_hook=event_hook, + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ) + assert explicitly_supported.event_hook == event_hook + + +@pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.post_call, + "post_call", + [GuardrailEventHooks.logging_only, GuardrailEventHooks.post_call], + ["logging_only", "post_call"], + Mode(tags={"enforce": "post_call"}, default="logging_only"), + Mode(tags={"enforce": ["logging_only", "post_call"]}), + Mode(tags={"audit": "logging_only"}, default="post_call"), + Mode(tags={}, default=["logging_only", "post_call"]), + ), +) +def test_framework_logging_only_does_not_allow_other_unsupported_modes( + event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, +) -> None: + with pytest.raises(ValueError, match=r"post_call.*not in the supported event hooks"): + _InheritedApplyGuardrail(event_hook=event_hook, supported_event_hooks=[GuardrailEventHooks.pre_call]) + + class TestLoggingOnlyApplyGuardrail: """LIT-4876 regression: a guardrail in mode logging_only that implements only apply_guardrail must still run against the logged request and response and record guardrail_information, instead of inheriting the CustomLogger no-op.""" + @pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.logging_only, + "logging_only", + [GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ["pre_call", "logging_only"], + Mode(tags={"audit": "logging_only"}, default="pre_call"), + Mode(tags={"audit": ["pre_call", "logging_only"]}), + Mode(tags={"enforce": "pre_call"}, default="logging_only"), + Mode(tags={}, default=["pre_call", "logging_only"]), + ), + ) + @pytest.mark.asyncio + async def test_content_filter_accepts_logging_only_and_records_detection( + self, event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode + ) -> None: + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks + + guardrail: Final = ContentFilterGuardrail( + guardrail_name="content-review", + event_hook=event_hook, + default_on=True, + blocked_words=[BlockedWord(keyword="hello", action=ContentFilterAction.BLOCK)], + ) + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert out_response is response + assert out_kwargs["messages"] == kwargs["messages"] + assert ( + out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"] + == "guardrail_intervened" + ) + @pytest.mark.asyncio async def test_runs_apply_guardrail_observe_only_and_records_verdict(self): guardrail = _ApplyOnlyObserver() @@ -2610,3 +2708,36 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: ) assert result is replacement + + @pytest.mark.asyncio + async def test_apply_guardrail_interface_modifies_deployment_response(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + class ReplacingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + assert input_type == "response" + return {**inputs, "texts": ["filtered response"]} + + guardrail = ReplacingGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.post_call, + ) + response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "original response"}}]) + request_data = {"guardrails": ["test-guardrail"]} + + result = await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, + response=response, + call_type=CallTypes.acompletion, + ) + + assert result is response + assert response.choices[0].message.content == "filtered response" + assert request_data == {"guardrails": ["test-guardrail"]} 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 e1bd20ece6f..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) 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_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 1778eca25ef..42d3df76902 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -9,9 +9,11 @@ import litellm from litellm.litellm_core_utils.exception_mapping_utils import ( ExceptionCheckers, _get_body_error_code, + _get_response_headers, exception_type, extract_and_raise_litellm_exception, ) +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.openai.common_utils import OpenAIError from litellm.types.utils import LlmProviders @@ -1254,3 +1256,156 @@ def test_handle_error_marks_only_a_status_code_it_never_received(): raise handler._handle_error(e=upstream, provider_config=None) assert received.value.status_code == 500 assert received.value.status_code_is_synthesized is False + + +def test_bedrock_500_preserves_provider_response_headers(): + """A Bedrock 5xx must keep x-amzn-RequestId so AWS support can trace it (LIT-5428).""" + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-map-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=500, + message=provider_response.text, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-map-500" + + +@pytest.mark.parametrize( + "custom_llm_provider, status_code, provider_message, expected_exception", + [ + ( + "bedrock_mantle", + 400, + ( + '{"error":{"code":"validation_error",' + '"message":"prompt tokens (1055489) exceed model maximum (1050000) for openai.gpt-5.6-sol",' + '"param":null,"type":"invalid_request_error"}}' + ), + litellm.ContextWindowExceededError, + ), + ( + "bedrock", + 400, + '{"message":"Input is too long for requested model."}', + litellm.ContextWindowExceededError, + ), + ( + "bedrock", + 400, + '{"message":"Could not process image"}', + litellm.InternalServerError, + ), + ], +) +def test_bedrock_classified_errors_preserve_provider_response_headers( + custom_llm_provider, status_code, provider_message, expected_exception +): + """Branches that classify a Bedrock error by its text must keep x-amzn-RequestId (LIT-5428).""" + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-classified"}, + text=provider_message, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message=provider_message, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(expected_exception) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider=custom_llm_provider, + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-classified" + + +@pytest.mark.parametrize( + "status_code, provider_message", + [ + (504, '{"message":"Gateway timeout"}'), + (408, '{"message":"Bedrock did not answer in time"}'), + (408, '{"message":"Connect timeout on endpoint URL"}'), + ], +) +def test_bedrock_timeout_mapping_preserves_provider_headers(status_code, provider_message): + """A mapped bedrock timeout keeps the upstream response, like every other mapped bedrock error. + + The proxy prefixes those headers on the way out, while retry and cooldown + logic still reads the raw retry-after off the response. + """ + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-timeout", "set-cookie": "session=attacker"}, + text=provider_message, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message=provider_message, + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.Timeout) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-timeout" + assert exc_info.value.headers is None + + +@pytest.mark.parametrize("status_code", [504, 408]) +def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code): + """Cooldown and retry timing read retry-after through _get_response_headers.""" + provider_response = httpx.Response( + status_code=status_code, + headers={"x-amzn-RequestId": "req-retry-after", "retry-after": "7"}, + text='{"message":"Bedrock did not answer in time"}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + original_exception = BedrockError( + status_code=status_code, + message='{"message":"Bedrock did not answer in time"}', + headers=provider_response.headers, + response=provider_response, + ) + + with pytest.raises(litellm.Timeout) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=original_exception, + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + exception_headers = _get_response_headers(original_exception=exc_info.value) + assert exception_headers is not None + assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7 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 0fdca755685..b6366bc803e 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -22,7 +22,13 @@ from litellm.litellm_core_utils.litellm_logging import ( _get_status_fields, set_callbacks, ) -from litellm.types.utils import ModelResponse, TextCompletionResponse +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import ( + CallTypes, + LiteLLMRealtimeStreamLoggingObject, + ModelResponse, + TextCompletionResponse, +) @pytest.fixture @@ -6393,6 +6399,90 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" +def _responses_ws_logging_obj() -> LitellmLogging: + return LitellmLogging( + model="gpt-4o", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="responses-ws-usage-test", + function_id="responses-ws-usage-test", + ) + + +def test_normalize_logging_result_extracts_usage_for_responses_websocket(monkeypatch): + """LIT-6512: native /v1/responses WebSocket sessions logged $0 spend because the usage + carried by stored response.completed events was never extracted. The session must cost + exactly what the same usage costs over HTTP /v1/responses, discounts included.""" + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.5}) + logging_obj = _responses_ws_logging_obj() + events = [ + {"type": "response.created", "response": {}}, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140}}, + }, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 60, "output_tokens": 10, "total_tokens": 70}}, + }, + ] + + normalized = logging_obj.normalize_logging_result(result=events) + + assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject) + assert normalized.usage.prompt_tokens == 160 + assert normalized.usage.completion_tokens == 50 + + ws_cost = litellm.completion_cost( + completion_response=normalized, + model="gpt-4o", + call_type=CallTypes.aresponses_websocket.value, + custom_llm_provider="openai", + ) + http_cost = litellm.completion_cost( + completion_response=ResponsesAPIResponse( + id="resp-6512", + created_at=1700000000, + output=[], + usage=ResponseAPIUsage(input_tokens=160, output_tokens=50, total_tokens=210), + ), + model="gpt-4o", + call_type=CallTypes.aresponses.value, + custom_llm_provider="openai", + ) + assert ws_cost > 0 + assert ws_cost == http_cost + + +def test_normalize_logging_result_bills_incomplete_responses_websocket_turns(): + """LIT-6512: a turn cut short by max_output_tokens ends in response.incomplete, which + OpenAI bills, so its usage counts toward the session like a completed turn.""" + events = [ + { + "type": "response.created", + "response": {"usage": {"input_tokens": 999, "output_tokens": 999, "total_tokens": 1998}}, + }, + { + "type": "response.incomplete", + "response": {"usage": {"input_tokens": 15, "output_tokens": 16, "total_tokens": 31}}, + }, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 40, "output_tokens": 4, "total_tokens": 44}}, + }, + {"type": "response.failed", "response": {"usage": None}}, + ] + + normalized = _responses_ws_logging_obj().normalize_logging_result(result=events) + + assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject) + assert normalized.usage.prompt_tokens == 55 + assert normalized.usage.completion_tokens == 20 + assert normalized.usage.total_tokens == 75 + + def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs).""" diff --git a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py index ca25ee80c23..d24e2b58db8 100644 --- a/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py +++ b/tests/test_litellm/litellm_core_utils/test_xai_oauth_routing.py @@ -1,6 +1,5 @@ -import litellm from litellm import LlmProviders from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -46,14 +45,6 @@ def test_xai_openai_compatible_provider_info(): assert dynamic_api_key == "api-key" -def test_xai_get_model_info_uses_xai_pricing_metadata(): - model_info = litellm.get_model_info("xai/grok-3-mini") - - assert model_info["litellm_provider"] == "xai" - assert model_info["key"] == "xai/grok-3-mini" - assert model_info["mode"] == "chat" - - def test_xai_validate_environment_reads_api_key(monkeypatch): monkeypatch.setenv("XAI_API_KEY", "api-key") diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 794613942a1..ae620fdd6dc 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -26,6 +26,30 @@ FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789" FAKE_AUTH_TOKEN = "sk-ant-aut01-fake-auth-token-for-testing-123456789" +@pytest.mark.parametrize( + "messages,system,expected", + [ + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_is_subagent=true;", True), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: =junk; cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: cc_version=; cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], "x-anthropic-billing-header: malformed", False), + ([{"content": "missing role"}], "x-anthropic-billing-header: cc_is_subagent=true;", False), + (["not-a-mapping"], "x-anthropic-billing-header: cc_is_subagent=true;", False), + ([{"role": "user", "content": "hi"}], ["not-a-mapping"], False), + ([{"role": "user", "content": "hi"}], None, False), + ], +) +def test_is_claude_code_one_shot_subagent_request(messages, system, expected): + from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request + + assert is_claude_code_one_shot_subagent_request( + messages=messages, + system=system, + tools=None, + user_agent="claude-cli/2.1.263 (external, cli)", + ) is expected + + class TestOptionallyHandleAnthropicOAuth: """Tests for optionally_handle_anthropic_oauth function.""" diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py index 2a44e77ce09..9bdc79919d2 100644 --- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py +++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py @@ -1,4 +1,3 @@ -import os from unittest.mock import MagicMock import httpx @@ -38,25 +37,6 @@ class TestAzureMAIImageGeneration: assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("flux.2-pro") assert not AzureFoundryMAIImageGenerationConfig.is_mai_model("MAI-DS-R1") - def test_mai_flash_and_2e_model_pricing_in_cost_map(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - - flash_info = litellm.get_model_info( - model="azure_ai/MAI-Image-2.5-Flash", - custom_llm_provider="azure_ai", - ) - assert flash_info["input_cost_per_token"] == 1.75e-06 - assert flash_info["input_cost_per_image_token"] == 1.75e-06 - assert flash_info["output_cost_per_image_token"] == 3.3e-05 - - image_2e_info = litellm.get_model_info( - model="azure_ai/MAI-Image-2e", - custom_llm_provider="azure_ai", - ) - assert image_2e_info["input_cost_per_token"] == 5e-06 - assert image_2e_info["output_cost_per_image_token"] == 1.95e-05 - def test_get_mai_image_generation_url(self): url = AzureFoundryMAIImageGenerationConfig.get_mai_image_generation_url( api_base="https://my-resource.services.ai.azure.com", diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py index f3618572622..1b2ca298694 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_fw_models_metadata.py @@ -12,112 +12,6 @@ from importlib.resources import files import pytest -FW_MODELS = { - "azure_ai/FW-Kimi-K2.5": { - "input_cost_per_token": 6.6e-07, - "output_cost_per_token": 3.3e-06, - "cache_read_input_token_cost": 1.1e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K2.6": { - "input_cost_per_token": 1.045e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 1.76e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K2.7-Code": { - "input_cost_per_token": 1.05e-06, - "output_cost_per_token": 4.4e-06, - "cache_read_input_token_cost": 2.1e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - "supports_vision": True, - }, - "azure_ai/FW-Kimi-K3": { - "input_cost_per_token": 3.3e-06, - "output_cost_per_token": 1.65e-05, - "cache_read_input_token_cost": 3.3e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - "supports_vision": True, - }, - "azure_ai/FW-Inkling": { - "input_cost_per_token": 1e-06, - "output_cost_per_token": 4.05e-06, - "cache_read_input_token_cost": 1.7e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 1048576, - }, - "azure_ai/FW-DeepSeek-V3.2": { - "input_cost_per_token": 6.2e-07, - "output_cost_per_token": 1.85e-06, - "cache_read_input_token_cost": 3.1e-07, - "max_input_tokens": 163840, - "max_output_tokens": 163840, - }, - "azure_ai/FW-DeepSeek-V4-Pro": { - "input_cost_per_token": 1.925e-06, - "output_cost_per_token": 3.828e-06, - "cache_read_input_token_cost": 1.65e-07, - "max_input_tokens": 1000000, - "max_output_tokens": 384000, - }, - "azure_ai/FW-MiniMax-M3": { - "input_cost_per_token": 3.3e-07, - "output_cost_per_token": 1.32e-06, - "cache_read_input_token_cost": 6.6e-08, - "max_input_tokens": 512000, - "max_output_tokens": 512000, - "supports_vision": True, - }, - "azure_ai/FW-MiniMax-M2.5": { - "input_cost_per_token": 3.3e-07, - "output_cost_per_token": 1.32e-06, - "cache_read_input_token_cost": 3.3e-08, - "max_input_tokens": 1000000, - "max_output_tokens": 1000000, - }, - "azure_ai/FW-Nemotron-3-Ultra-NVFP4": { - "input_cost_per_token": 6e-07, - "output_cost_per_token": 2.4e-06, - "cache_read_input_token_cost": 1.19e-07, - "max_input_tokens": 262144, - "max_output_tokens": 262144, - }, - "azure_ai/FW-GLM-5.2-Fast": { - "input_cost_per_token": 2.1e-06, - "output_cost_per_token": 6.6e-06, - "cache_read_input_token_cost": 2.1e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5.2": { - "input_cost_per_token": 1.54e-06, - "output_cost_per_token": 4.84e-06, - "cache_read_input_token_cost": 1.5e-07, - "max_input_tokens": 1048576, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5.1": { - "input_cost_per_token": 1.54e-06, - "output_cost_per_token": 4.84e-06, - "cache_read_input_token_cost": 2.86e-07, - "max_input_tokens": 202800, - "max_output_tokens": 131072, - }, - "azure_ai/FW-GLM-5": { - "input_cost_per_token": 1.1e-06, - "output_cost_per_token": 3.52e-06, - "cache_read_input_token_cost": 2.2e-07, - "max_input_tokens": 200000, - "max_output_tokens": 128000, - }, -} - @pytest.fixture(scope="module") def use_local_model_cost_map(): @@ -144,28 +38,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -@pytest.mark.parametrize("model_key,expected", list(FW_MODELS.items())) -def test_azure_ai_fw_model_info(use_local_model_cost_map, model_key, expected): - model_info = use_local_model_cost_map.get_model_info(model=model_key) - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["input_cost_per_token"] == pytest.approx(expected["input_cost_per_token"]) - assert model_info["output_cost_per_token"] == pytest.approx(expected["output_cost_per_token"]) - assert model_info["cache_read_input_token_cost"] == pytest.approx( - expected["cache_read_input_token_cost"] - ) - assert model_info["max_input_tokens"] == expected["max_input_tokens"] - assert model_info["max_output_tokens"] == expected["max_output_tokens"] - assert model_info["max_tokens"] == expected["max_output_tokens"] - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_prompt_caching"] is True - if expected.get("supports_vision"): - assert model_info["supports_vision"] is True - - @pytest.mark.parametrize( "model_name,expected_prompt,expected_completion", [ @@ -197,22 +69,6 @@ def test_azure_ai_fw_cost_per_token( assert completion_cost == pytest.approx(expected_completion) -def test_azure_ai_fw_nemotron_lightning_model_info(use_local_model_cost_map): - model_info = use_local_model_cost_map.get_model_info(model="azure_ai/FW-Nemotron-Lightning-3.5-30B-A3B") - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["input_cost_per_token"] == pytest.approx(6e-08) - assert model_info["output_cost_per_token"] == pytest.approx(2.2e-07) - assert model_info["cache_read_input_token_cost"] == pytest.approx(1e-08) - assert model_info["max_input_tokens"] == 262144 - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_prompt_caching"] is True - assert model_info["supports_vision"] is False - - def test_azure_ai_fw_nemotron_lightning_supports_tool_choice(use_local_model_cost_map): from litellm.llms.azure_ai.chat.transformation import AzureAIStudioConfig diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py index 812b9288ca8..cbcc2a94043 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_kimi_k26_metadata.py @@ -33,33 +33,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -def test_azure_ai_kimi_k26_model_info(use_local_model_cost_map): - model_info = use_local_model_cost_map.get_model_info(model="azure_ai/kimi-k2.6") - - assert model_info["litellm_provider"] == "azure_ai" - assert model_info["mode"] == "chat" - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - assert model_info["max_tokens"] == 262144 - assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07) - assert model_info["output_cost_per_token"] == pytest.approx(4e-06) - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - -def test_azure_ai_kimi_k26_raw_model_cost_entry(use_local_model_cost_map): - model_info = use_local_model_cost_map.model_cost["azure_ai/kimi-k2.6"] - - assert model_info["supported_modalities"] == ["text", "image"] - assert model_info["supported_output_modalities"] == ["text"] - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - def test_azure_ai_kimi_k26_cost_per_token(use_local_model_cost_map): from litellm.llms.azure_ai.cost_calculator import cost_per_token from litellm.types.utils import Usage diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index aba51689094..c2c448cd7e2 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -177,3 +177,16 @@ def test_guardrail_config_flows_to_headers_not_request_body(model): assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "ff6ujrregl1q" assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT" assert headers["X-Amzn-Bedrock-Trace"] == "DISABLED" + + +def test_get_error_class_preserves_provider_headers(): + """The invoke handler path hands real provider headers to get_error_class (LIT-5428).""" + error = AmazonInvokeConfig().get_error_class( + error_message="Amazon Bedrock is unable to process your request.", + status_code=500, + headers={"x-amzn-RequestId": "req-invoke-500"}, + ) + + assert isinstance(error, BedrockError) + assert error.headers == {"x-amzn-RequestId": "req-invoke-500"} + assert error.response.headers["x-amzn-requestid"] == "req-invoke-500" diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index cb05cdb9451..f0e361ceb88 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2,6 +2,7 @@ import asyncio import json import os +import httpx import pytest from fastapi.testclient import TestClient @@ -6039,6 +6040,8 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}} class MockResponse: + headers = httpx.Headers({"x-amzn-RequestId": "req-parse-failure"}) + def json(self): return leaky_body @@ -6067,6 +6070,7 @@ def test_transform_response_does_not_leak_body_on_parse_failure(): msg = str(exc_info.value) assert "secret content" not in msg assert "Error converting to valid response block" in msg + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-parse-failure" def test_converse_drops_sampling_params_for_models_that_removed_them(): diff --git a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py index 4bef59842f1..d0adabe7b4e 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_invoke_handler.py @@ -496,3 +496,171 @@ async def test_async_invoke_streaming_forwards_bedrock_response_headers(): assert stream._hidden_params["additional_headers"]["llm_provider-x-amzn-requestid"] == "req-987" + +def _bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_invoke_streaming_error_forwards_bedrock_response_headers(): + error_response = _bedrock_stream_error_response(500, "req-stream-err-1") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-1" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_error_forwards_bedrock_response_headers(): + error_response = _bedrock_stream_error_response(500, "req-stream-err-2") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-stream-err-2" + + +def _unread_bedrock_stream_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + stream=httpx.ByteStream(b'{"message":"Amazon Bedrock is unable to process your request."}'), + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_invoke_streaming_error_forwards_headers_when_body_was_never_read(): + """A retried streamed request raises HTTPStatusError over a body nobody read, so + reading it for the error message throws and loses the request id (LIT-5428).""" + error_response = _unread_bedrock_stream_error_response(500, "req-unread-sync") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-sync" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_error_forwards_headers_when_body_was_never_read(): + error_response = _unread_bedrock_stream_error_response(500, "req-unread-async") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-unread-async" + + +def test_invoke_streaming_non_200_forwards_bedrock_response_headers(): + """A caller-supplied client that returns a failure instead of raising still reaches the + provider's headers, and reading the streamed body for the message must not throw (LIT-5428).""" + error_response = _unread_bedrock_stream_error_response(500, "req-non200-sync") + client = HTTPHandler() + client.post = MagicMock(return_value=error_response) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-sync" + + +@pytest.mark.asyncio +async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers(): + error_response = _unread_bedrock_stream_error_response(500, "req-non200-async") + client = AsyncHTTPHandler() + client.post = AsyncMock(return_value=error_response) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-non200-async" diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 9302dc01abe..3f03305423a 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -614,3 +614,260 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] assert "ASIABATCHSIGNROLE" in authorization assert signed_data == b'{"jobName": "litellm-batch-job"}' + + +# --------------------------------------------------------------------------- # +# Provider error headers (LIT-5428) # +# --------------------------------------------------------------------------- # + + +def _bedrock_chat_error_configs(): + from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + from litellm.llms.bedrock.chat.invoke_agent.transformation import AmazonInvokeAgentConfig + from litellm.llms.bedrock.chat.invoke_transformations.amazon_moonshot_transformation import ( + AmazonMoonshotConfig, + ) + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, + ) + from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig, + ) + + return [ + AmazonInvokeConfig, + AmazonConverseConfig, + AmazonMoonshotConfig, + AmazonBedrockOpenAIConfig, + AmazonAgentCoreConfig, + AmazonInvokeAgentConfig, + ] + + +@pytest.mark.parametrize("config", _bedrock_chat_error_configs()) +def test_bedrock_chat_get_error_class_keeps_provider_headers(config): + """Every Bedrock chat route must carry x-amzn-RequestId out to the caller (LIT-5428). + + A config that drops the headers it is handed shadows the fix for its own models. + """ + error = config().get_error_class( + error_message="Amazon Bedrock is unable to process your request.", + status_code=500, + headers={"x-amzn-RequestId": "req-chat-500"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-chat-500" + + +def test_error_response_text_reads_a_read_response(): + import httpx + + from litellm.llms.bedrock.common_utils import error_response_text + + response = httpx.Response(status_code=500, text="Amazon Bedrock is unable to process your request.") + + assert error_response_text(response) == "Amazon Bedrock is unable to process your request." + + +def test_error_response_text_falls_back_when_a_streamed_response_was_never_read(): + """A retried streamed request raises HTTPStatusError over an unread body; reading it + throws ResponseNotRead and would lose the status and headers this fix preserves.""" + import httpx + + from litellm.llms.bedrock.common_utils import error_response_text + + request = httpx.Request(method="POST", url="https://bedrock-runtime.amazonaws.com") + response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-unread-500"}, + stream=httpx.ByteStream(b"never read"), + request=request, + ) + + with pytest.raises(httpx.ResponseNotRead): + _ = response.text + + assert error_response_text(response) == "Internal Server Error" + + +def test_bedrock_error_skips_header_values_httpx_cannot_carry(): + """The shared HTTP handler copies an arbitrary exception's header values in verbatim, + so a non-str value must not take down the whole error (LIT-5428).""" + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="boom", + headers={"x-amzn-RequestId": "req-mixed-500", "x-retry-count": 3, "x-nothing": None}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-mixed-500" + assert "x-retry-count" not in error.response.headers + assert isinstance(error.response, httpx.Response) + + +def test_bedrock_error_keeps_duplicate_httpx_header_values(): + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="boom", + headers=httpx.Headers([("x-amzn-RequestId", "req-dup-500"), ("set-cookie", "a=1"), ("set-cookie", "b=2")]), + ) + + assert error.response.headers.get_list("set-cookie") == ["a=1", "b=2"] + + +def _bedrock_httpx_status_error_sites(): + """Every `except httpx.HTTPStatusError as err` that raises a BedrockError, across bedrock.""" + import ast + import pathlib + + sites = [] + for path in sorted(pathlib.Path("litellm/llms/bedrock").rglob("*.py")): + tree = ast.parse(path.read_text()) + for handler in (n for n in ast.walk(tree) if isinstance(n, ast.ExceptHandler)): + caught = ast.unparse(handler.type) if handler.type is not None else "" + if "HTTPStatusError" not in caught or handler.name is None: + continue + for call in ( + n + for n in ast.walk(handler) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "BedrockError" + ): + sites.append((str(path), call.lineno, handler.name, {k.arg for k in call.keywords})) + return sites + + +def test_every_bedrock_httpx_status_error_site_keeps_provider_headers(): + """A raise site holding the provider's failed response must hand its headers on (LIT-5428). + + These sites are the only place x-amzn-RequestId still exists; a site that drops it + silently shadows the fix for that whole surface. + """ + sites = _bedrock_httpx_status_error_sites() + + assert len(sites) >= 12 + dropped = [f"{path}:{lineno}" for path, lineno, _, kwargs in sites if "headers" not in kwargs] + assert dropped == [] + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_bedrock_embedding_call_keeps_provider_headers(is_async): + """The embeddings surface raises from the same shape as chat and lost the same header.""" + import httpx + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + + failure = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-embed-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + class _SyncUpstream(HTTPHandler): + def post(self, *args, **kwargs): + return failure + + class _AsyncUpstream(AsyncHTTPHandler): + async def post(self, *args, **kwargs): + return failure + + async def _drive(): + embedding = BedrockEmbedding() + kwargs = dict( + timeout=None, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/", + headers={}, + data={}, + ) + if is_async: + return await embedding._make_async_call(client=_AsyncUpstream(), **kwargs) + return embedding._make_sync_call(client=_SyncUpstream(), **kwargs) + + with pytest.raises(BedrockError) as exc_info: + await _drive() + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-embed-500" + + +def _bedrock_mantle_error_configs(): + from litellm.llms.bedrock_mantle.chat.transformation import BedrockMantleChatConfig + from litellm.llms.bedrock_mantle.responses.transformation import BedrockMantleResponsesAPIConfig + + return [BedrockMantleChatConfig, BedrockMantleResponsesAPIConfig] + + +@pytest.mark.parametrize("config", _bedrock_mantle_error_configs()) +def test_bedrock_mantle_get_error_class_keeps_provider_headers(config): + """bedrock_mantle rides the OpenAI-compatible surfaces, whose errors drop the headers. + + A chat request for a responses-API model is bridged onto the responses config, so + fixing only the chat one leaves the model the customer actually calls uncovered. + """ + error = config().get_error_class( + error_message="prompt tokens exceed model maximum", + status_code=400, + headers={"x-amzn-RequestId": "req-mantle-400"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-mantle-400" + + +def _bedrock_configs_with_get_error_class(): + import importlib + import inspect + import pathlib + + import litellm + + llms_root = pathlib.Path(inspect.getfile(litellm)).parent / "llms" + configs = [] + for package in ("bedrock", "bedrock_mantle"): + for path in sorted((llms_root / package).rglob("*.py")): + module_name = "litellm.llms." + ".".join(path.relative_to(llms_root).with_suffix("").parts) + module = importlib.import_module(module_name) + for name, obj in vars(module).items(): + if not inspect.isclass(obj) or obj.__module__ != module_name: + continue + if getattr(obj, "get_error_class", None) is None: + continue + configs.append(pytest.param(obj, id=f"{module_name}.{name}")) + return configs + + +@pytest.mark.parametrize("config", _bedrock_configs_with_get_error_class()) +def test_every_bedrock_config_get_error_class_keeps_provider_headers(config): + """Every bedrock surface must classify errors through BedrockError, not a header-dropping base. + + A config that inherits get_error_class from a provider-agnostic base builds a blank + response, so the request id is gone before the proxy ever reads it. + """ + try: + instance = config() + except Exception: + instance = config.__new__(config) + + try: + error = instance.get_error_class( + error_message="boom", + status_code=500, + headers={"x-amzn-RequestId": "req-audit-500"}, + ) + except Exception as raised: # some bases raise the exception instead of returning it + error = raised + + assert error.response.headers["x-amzn-requestid"] == "req-audit-500" + + +def test_bedrock_get_error_class_audit_covers_every_surface(): + assert len(_bedrock_configs_with_get_error_class()) >= 30 diff --git a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py index 5f12ae8566c..fda3c8ceb8f 100644 --- a/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py +++ b/tests/test_litellm/llms/bedrock/test_cross_region_inference_profile_mapping.py @@ -1,8 +1,5 @@ """Test Bedrock cross-region inference profile model mapping""" -import json -from functools import lru_cache -from pathlib import Path from typing import NamedTuple import pytest @@ -102,13 +99,6 @@ GPT_5_6_PROFILES = [ ] -@lru_cache(maxsize=1) -def _packaged_cost_map(): - """The map litellm actually resolves against, for fields ModelInfoBase drops.""" - path = Path(litellm.__file__).parent / "model_prices_and_context_window_backup.json" - return json.loads(path.read_text()) - - def _bedrock_response(model, usage): return ModelResponse( id="test", @@ -126,17 +116,6 @@ def _bedrock_response(model, usage): ) -def test_bedrock_cross_region_inference_profile_mapping(): - """Test that bedrock cross-region inference profile model is mapped""" - model = "bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0" - - model_info = _get_model_info_helper(model=model, custom_llm_provider="bedrock") - - assert model_info is not None - assert model_info["litellm_provider"] == "bedrock" - assert model_info["input_cost_per_token"] == 8e-07 - - def test_proxy_cost_calculation_scenario(): """Test exact GitHub issue scenario: proxy cost calculation""" model = "litellm_proxy/bedrock/us.anthropic.claude-3-5-haiku-20241022-v1:0" @@ -176,38 +155,6 @@ def test_bedrock_gpt_5_6_profiles_route_to_converse(profile, local_model_cost_ma assert BedrockModelInfo.get_bedrock_route(f"bedrock/{profile.model_id}") == "converse" -@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) -def test_bedrock_gpt_5_6_published_rates(profile, local_model_cost_map): - """Geo and Global profiles carry their own published rates, per context tier.""" - model_info = _get_model_info_helper( - model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" - ) - - assert model_info["litellm_provider"] == "bedrock_converse" - assert model_info["mode"] == "chat" - assert model_info["max_input_tokens"] == 1000000 - assert model_info["input_cost_per_token"] == profile.input_cost - assert ( - model_info["input_cost_per_token_above_272k_tokens"] - == profile.input_cost_above_272k - ) - assert model_info["output_cost_per_token"] == profile.output_cost - assert ( - model_info["output_cost_per_token_above_272k_tokens"] - == profile.output_cost_above_272k - ) - assert model_info["cache_creation_input_token_cost"] == profile.cache_write - assert ( - model_info["cache_creation_input_token_cost_above_272k_tokens"] - == profile.cache_write_above_272k - ) - assert model_info["cache_read_input_token_cost"] == profile.cache_read - assert ( - model_info["cache_read_input_token_cost_above_272k_tokens"] - == profile.cache_read_above_272k - ) - - def test_bedrock_gpt_5_6_above_272k_tier_applies_to_cost(local_model_cost_map): """A prompt over 272K tokens is billed at the long-context rate, not the base rate.""" response = _bedrock_response( @@ -267,31 +214,6 @@ def test_bedrock_gpt_5_6_bills_cache_write_tokens(local_model_cost_map): assert cost == pytest.approx(expected, rel=1e-9) -@pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) -def test_bedrock_gpt_5_6_advertises_only_converse_supported_features( - profile, local_model_cost_map -): - model_info = _get_model_info_helper( - model=f"bedrock/{profile.model_id}", custom_llm_provider="bedrock" - ) - - assert model_info["supports_function_calling"] is True - assert model_info["supports_tool_choice"] is True - assert model_info["supports_vision"] is True - - # Bedrock rejects an explicit cachePoint block for these models, so the flag that - # offers caller-driven caching stays off even though the cache rates are declared. - assert not model_info.get("supports_prompt_caching") - - # ModelInfoBase drops these two, so they are read from the map litellm resolves. - raw = _packaged_cost_map()[profile.model_id] - assert raw["supported_modalities"] == ["text", "image"] - assert raw["supported_output_modalities"] == ["text"] - # No bedrock_converse entry declares supported_endpoints; these models are reachable - # on chat completions and on the Responses API without it. - assert "supported_endpoints" not in raw - - @pytest.mark.parametrize("profile", GPT_5_6_PROFILES, ids=lambda p: p.model_id) def test_bedrock_gpt_5_6_offers_tools_and_reasoning_effort_but_not_thinking(profile, local_model_cost_map): """GPT-5.x on Converse maps reasoning_effort to reasoning.effort, so reasoning_effort diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py index 1f23d39c631..9457d5faaff 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py @@ -8,9 +8,7 @@ gate, the URL construction for both paths, and the shared Bearer auth. """ import copy -import json import logging -from pathlib import Path import pytest from botocore.exceptions import ( @@ -159,7 +157,6 @@ class TestBedrockMantleResponsesURL: assert url == "https://bedrock-mantle.us-east-2.api.aws/v1/responses" assert url.count("/responses") == 1 - def test_url_aws_region_name_overrides_stale_api_base(self, monkeypatch): monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) @@ -1777,53 +1774,6 @@ class TestBedrockMantleResponsesSigV4: class TestBedrockMantleResponsesPricing: - def test_gpt_5_5_pricing_and_mode(self, local_cost_map): - info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.5") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(5.5e-06) - assert info["output_cost_per_token"] == pytest.approx(3.3e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(5.5e-07) - assert info["max_input_tokens"] == 1050000 - - def test_gpt_5_4_pricing_and_mode(self, local_cost_map): - info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.4") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(2.75e-06) - assert info["output_cost_per_token"] == pytest.approx(1.65e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(2.75e-07) - assert info["max_input_tokens"] == 1050000 - - def test_gpt_5_6_cyber_pricing_and_mode(self, local_cost_map): - info = litellm.get_model_info("bedrock_mantle/openai.gpt-5.6-cyber") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(1.375e-05) - assert info["cache_creation_input_token_cost"] == pytest.approx(1.71875e-05) - assert info["cache_read_input_token_cost"] == pytest.approx(1.375e-06) - assert info["output_cost_per_token"] == pytest.approx(8.25e-05) - assert info["max_input_tokens"] == 272000 - - @pytest.mark.parametrize( - "model, input_cost, cache_creation_cost, cache_read_cost, output_cost", - [ - ("openai.gpt-5.6-sol", 5.5e-06, 6.875e-06, 5.5e-07, 3.3e-05), - ("openai.gpt-5.6-terra", 2.2e-06, 2.75e-06, 2.2e-07, 1.32e-05), - ("openai.gpt-5.6-luna", 2.2e-07, 2.75e-07, 2.2e-08, 1.32e-06), - ], - ) - def test_gpt_5_6_pricing_and_mode( - self, local_cost_map, model, input_cost, cache_creation_cost, cache_read_cost, output_cost - ): - info = litellm.get_model_info(f"bedrock_mantle/{model}") - assert info["mode"] == "responses" - assert info["input_cost_per_token"] == pytest.approx(input_cost) - assert info["cache_creation_input_token_cost"] == pytest.approx(cache_creation_cost) - assert info["cache_read_input_token_cost"] == pytest.approx(cache_read_cost) - assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == 1050000 - assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(input_cost * 2) - assert info["cache_creation_input_token_cost_above_272k_tokens"] == pytest.approx(cache_creation_cost * 2) - assert info["cache_read_input_token_cost_above_272k_tokens"] == pytest.approx(cache_read_cost * 2) - assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(output_cost * 1.5) @pytest.mark.parametrize( "model, input_cost, output_cost", @@ -1861,58 +1811,3 @@ class TestBedrockMantleResponsesPricing: def test_models_registered(self, local_cost_map): assert "bedrock_mantle/openai.gpt-5.5" in litellm.bedrock_mantle_models assert "bedrock_mantle/openai.gpt-5.4" in litellm.bedrock_mantle_models - - -def _repo_cost_map(map_name: str) -> dict[str, dict[str, object]]: - repo_root = Path(__file__).resolve().parents[4] - paths = { - "root": repo_root / "model_prices_and_context_window.json", - "bundled_backup": repo_root / "litellm" / "model_prices_and_context_window_backup.json", - } - return json.loads(paths[map_name].read_text()) - - -class TestMantleGptRegistryEntries: - """Locks the OpenAI GPT entries to Bedrock Mantle's live behavior. - - Mantle enforces a 1,050,000-token prompt maximum for gpt-5.6 sol/terra/luna - and for gpt-5.5 and gpt-5.4 (oversize requests 400 with "prompt tokens (N) - exceed model maximum (1050000)", and a 1,030,590-token request completes - on every one of them), while the AWS model cards still quote 272K for - gpt-5.5 and gpt-5.4. mode must stay "responses": Mantle's native - /v1/chat/completions rejects function tools unless reasoning_effort is - "none", so chat traffic has to keep bridging to the Responses API - (see the responses_api_bridge tests above). - """ - - @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) - @pytest.mark.parametrize( - "key", - ( - "bedrock_mantle/openai.gpt-5.6-sol", - "bedrock_mantle/openai.gpt-5.6-terra", - "bedrock_mantle/openai.gpt-5.6-luna", - ), - ) - def test_entry_matches_mantle_enforced_limits(self, map_name, key): - entry = _repo_cost_map(map_name)[key] - assert entry["max_input_tokens"] == 1050000 - assert entry["max_output_tokens"] == 128000 - assert entry["mode"] == "responses" - assert entry["use_openai_responses_path"] is True - assert entry["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] - - @pytest.mark.parametrize("map_name", ("root", "bundled_backup")) - @pytest.mark.parametrize( - "key", - ( - "bedrock_mantle/openai.gpt-5.5", - "bedrock_mantle/openai.gpt-5.4", - ), - ) - def test_gpt_55_and_54_entries_match_mantle_enforced_limits(self, map_name, key): - entry = _repo_cost_map(map_name)[key] - assert entry["max_input_tokens"] == 1050000 - assert entry["max_output_tokens"] == 128000 - assert entry["mode"] == "responses" - assert entry["use_openai_responses_path"] is True diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index b88c27e64b9..1be94d4daa2 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -684,40 +684,6 @@ class TestBedrockMantleProviderResolution: class TestBedrockMantlePricing: """Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing.""" - def test_gpt_oss_120b_pricing(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - # Bedrock pricing: $0.15/M input, $0.60/M output - assert info["input_cost_per_token"] == pytest.approx(1.5e-7) - assert info["output_cost_per_token"] == pytest.approx(6e-7) - - def test_gpt_oss_20b_pricing(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-20b") - # Bedrock pricing: $0.075/M input, $0.30/M output - assert info["input_cost_per_token"] == pytest.approx(7.5e-8) - assert info["output_cost_per_token"] == pytest.approx(3e-7) - - def test_pricing_significantly_cheaper_than_openai_native(self, monkeypatch): - """ - Verify Bedrock Mantle pricing is cheaper than OpenAI's direct API pricing. - This is the core issue the provider addition fixes — previously users were being - billed at OpenAI rates instead of the cheaper Bedrock rates. - """ - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - bedrock_info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - # OpenAI direct pricing for gpt-oss-120b is ~$0.039/M input, $0.190/M output - # Bedrock should be cheaper at $0.15/M input and $0.60/M output... wait - # Actually, Bedrock ADDS value not reduces cost vs OpenAI direct for these models. - # The key fix is that we now use Bedrock-specific prices instead of mapping to - # some unrelated OpenAI model (like gpt-4) pricing. - # Just validate the pricing is as expected from AWS docs. - assert bedrock_info["input_cost_per_token"] == pytest.approx(1.5e-7) - assert bedrock_info["output_cost_per_token"] == pytest.approx(6e-7) - def test_safeguard_models_have_larger_output_tokens(self, monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") litellm.add_known_models() @@ -727,49 +693,6 @@ class TestBedrockMantlePricing: ) assert info_safeguard["max_output_tokens"] > info_120b["max_output_tokens"] - def test_reasoning_support(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - assert info.get("supports_reasoning") is True - - def test_context_window(self, monkeypatch): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true") - litellm.add_known_models() - info = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b") - assert info["max_input_tokens"] == 131072 - - -@pytest.mark.parametrize( - "model_id,input_cost,output_cost,max_tokens", - [ - ("google.gemma-4-31b", 1.4e-07, 4e-07, 256000), - ("google.gemma-4-26b-a4b", 1.3e-07, 4e-07, 256000), - ("google.gemma-4-e2b", 4e-08, 8e-08, 128000), - ], -) -def test_gemma_4_bedrock_mantle_model_metadata( - local_cost_map, model_id, input_cost, output_cost, max_tokens -): - full_model_name = f"bedrock_mantle/{model_id}" - info = litellm.get_model_info(full_model_name) - - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == pytest.approx(input_cost) - assert info["output_cost_per_token"] == pytest.approx(output_cost) - assert info["max_input_tokens"] == max_tokens - assert info["max_output_tokens"] == max_tokens - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert ( - litellm.supports_parallel_function_calling( - model=full_model_name, custom_llm_provider="bedrock_mantle" - ) - is False - ) - @pytest.mark.parametrize( "model_id", diff --git a/tests/test_litellm/llms/chat/test_converse_handler.py b/tests/test_litellm/llms/chat/test_converse_handler.py index ca79c8d7025..12b5f03aedc 100644 --- a/tests/test_litellm/llms/chat/test_converse_handler.py +++ b/tests/test_litellm/llms/chat/test_converse_handler.py @@ -308,3 +308,64 @@ def test_completion_plumbs_stream_chunk_size_through_converse(): stream_chunk_size=2048, ) iter_bytes_spy.assert_called_once_with(chunk_size=2048) + + +def _bedrock_error_response(status_code: int, request_id: str) -> httpx.Response: + return httpx.Response( + status_code=status_code, + headers={ + "x-amzn-RequestId": request_id, + "x-amzn-ErrorType": "InternalServerException", + }, + text=json.dumps({"message": "Amazon Bedrock is unable to process your request."}), + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + +def test_converse_completion_error_forwards_bedrock_response_headers(): + error_response = _bedrock_error_response(500, "req-err-123") + client = HTTPHandler() + client.post = MagicMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + litellm.completion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-123" + + +@pytest.mark.asyncio +async def test_async_converse_completion_error_forwards_bedrock_response_headers(): + error_response = _bedrock_error_response(500, "req-err-456") + client = AsyncHTTPHandler() + client.post = AsyncMock( + side_effect=httpx.HTTPStatusError( + "server error", + request=error_response.request, + response=error_response, + ) + ) + + with pytest.raises(litellm.ServiceUnavailableError) as exc_info: + await litellm.acompletion( + model="bedrock/converse/anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + client=client, + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + assert exc_info.value.response.headers["x-amzn-requestid"] == "req-err-456" 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 e16855da8cb..98a5f4b2db5 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 @@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams @@ -37,6 +38,69 @@ from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, Trans _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +OCR_RESPONSE = { + "pages": [{"index": 0, "markdown": "OCR output", "images": []}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, +} + + +def _ocr_sync_client() -> HTTPHandler: + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE))) + return client + + +def _ocr_async_client() -> AsyncHTTPHandler: + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient( + transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE)) + ) + return client + + +def test_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = BaseLLMHTTPHandler().ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_sync_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + + +@pytest.mark.asyncio +async def test_async_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = await BaseLLMHTTPHandler().async_ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_async_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index e6fe01be4ba..d4ef4282b27 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm -from litellm import get_model_info, supports_reasoning, supports_vision +from litellm import supports_reasoning, supports_vision from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id @@ -16,17 +16,6 @@ from litellm.types.utils import ( ) -@pytest.fixture(autouse=True) -def force_local_model_cost(monkeypatch): - """Force local model cost map usage for all tests in this file.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - # Refresh model_cost from local map - import litellm - from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map - - litellm.model_cost = get_model_cost_map(url=litellm.model_cost_map_url) - - def test_validate_environment_sets_session_affinity_from_litellm_session_id(): config = FireworksAIConfig() @@ -404,15 +393,6 @@ def test_get_supported_openai_params_parallel_tool_calls_without_tool_choice( assert "tool_choice" not in supported_params -def test_get_model_info_respects_explicit_fireworks_capabilities(): - """Test that get_model_info preserves explicit capability flags from the model map.""" - model_info = get_model_info("fireworks_ai/accounts/fireworks/models/glm-5p1") - - assert model_info["supports_function_calling"] is True - assert model_info["supports_reasoning"] is True - assert model_info["supports_tool_choice"] is True - - def test_get_provider_info_omits_false_supports_reasoning(monkeypatch): """Test that Fireworks only overrides supports_reasoning for supported models.""" config = FireworksAIConfig() diff --git a/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py new file mode 100644 index 00000000000..b9408d44e9a --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/responses/test_fireworks_ai_responses_transformation.py @@ -0,0 +1,455 @@ +import json +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, TypedDict, cast +from unittest.mock import MagicMock, patch +from urllib.parse import quote + +import httpx +import pytest +from openai.types.responses import ( + EasyInputMessage, + ResponseFunctionToolCall, + ResponseOutputMessage, + ResponseOutputText, + ResponseReasoningItem, +) +from openai.types.responses.response_input_param import FunctionCallOutput +from openai.types.responses.tool_param import Mcp +from typing_extensions import ReadOnly + +import litellm +from litellm.llms.fireworks_ai.responses.transformation import FireworksAIResponsesAPIConfig +from litellm.responses.file_search.emulated_handler import should_use_emulated_file_search +from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponseInputParam, ResponsesAPIResponse +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +FIREWORKS_RESPONSES_URL: Final = "https://api.fireworks.ai/inference/v1/responses" +HTTPX_CLIENT_FACTORY: Final = "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client" +NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) +NO_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _GeneratedSessionMetadata(TypedDict): + litellm_session_id_generated: ReadOnly[bool] + + +def _fireworks_response(model: str) -> Mapping[str, object]: + return ResponsesAPIResponse( + id="resp_0e946f2d46bf4b49bf8b29ff78083583", + object="response", + created_at=1788550000, + model=model, + status="completed", + output=( + ResponseReasoningItem(id="rs_1", summary=(), type="reasoning"), + ResponseOutputMessage( + id="msg_1", + status="completed", + role="assistant", + type="message", + content=(ResponseOutputText(type="output_text", text="Paris is clear and 21C.", annotations=()),), + ), + ResponseFunctionToolCall( + id="fc_1", + call_id="call_abc123", + name="get_weather", + arguments='{"city": "Paris"}', + status="completed", + type="function_call", + ), + ), + usage=ResponseAPIUsage( + input_tokens=179, + output_tokens=100, + total_tokens=279, + input_tokens_details=InputTokensDetails(cached_tokens=0), + ), + ).model_dump(mode="json", exclude_none=True) + + +def _mock_http_client(response_body: Mapping[str, object]) -> MagicMock: + client: Final = MagicMock() + response: Final = MagicMock() + response.status_code = 200 + response.headers = httpx.Headers((("content-type", "application/json"),)) + response.json.return_value = response_body + response.text = json.dumps(response_body) + client.post.return_value = response + return client + + +def _sent_request(client: MagicMock) -> tuple[str, Mapping[str, str], Mapping[str, object]]: + kwargs: Final = client.post.call_args.kwargs + body: Final = kwargs["json"] if "json" in kwargs else json.loads(kwargs["data"]) + return kwargs["url"], kwargs["headers"], body + + +@pytest.fixture(autouse=True) +def fireworks_env(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "FIREWORKS_API_KEY", + "FIREWORKS_AI_API_KEY", + "FIREWORKSAI_API_KEY", + "FIREWORKS_AI_TOKEN", + "FIREWORKS_API_BASE", + ): + monkeypatch.delenv(name, raising=False) + + +def test_fireworks_ai_provider_config_registration() -> None: + config: Final = ProviderConfigManager.get_provider_responses_api_config( + model="accounts/fireworks/models/kimi-k3", provider=LlmProviders.FIREWORKS_AI + ) + assert isinstance(config, FireworksAIResponsesAPIConfig) + assert config.custom_llm_provider == LlmProviders.FIREWORKS_AI + + +def test_responses_call_hits_native_endpoint_with_mcp_tool_untouched() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + mcp_tool: Final[Mcp] = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/mcp", + "require_approval": "never", + } + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + response: Final = litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input="What is litellm?", + tools=[mcp_tool], # mutable-ok: the Responses API takes tools as a JSON list + api_key="fw-test-key", + ) + url, headers, body = _sent_request(client) + assert url == FIREWORKS_RESPONSES_URL + assert headers["Authorization"] == "Bearer fw-test-key" + assert body["model"] == "accounts/fireworks/models/kimi-k3" + assert tuple(body["tools"]) == (mcp_tool,) + assert "messages" not in body + assert isinstance(response, ResponsesAPIResponse) + function_calls: Final = tuple(item for item in response.output if getattr(item, "type", None) == "function_call") + assert getattr(function_calls[0], "call_id", None) == "call_abc123" + + +def test_responses_call_expands_bare_model_name_to_fireworks_resource() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/glm-5p3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/glm-5p3", input="hi", api_key="fw-test-key") + _, _, body = _sent_request(client) + assert body["model"] == "accounts/fireworks/models/glm-5p3" + + +def test_responses_call_forwards_previous_response_id_and_store() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + tool_output: Final[FunctionCallOutput] = { + "type": "function_call_output", + "call_id": "call_abc123", + "output": "{}", + } + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input=[tool_output], # mutable-ok: the Responses API takes input items as a JSON list + previous_response_id="resp_0e946f2d46bf4b49bf8b29ff78083583", + store=True, + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert body["previous_response_id"] == "resp_0e946f2d46bf4b49bf8b29ff78083583" + assert body["store"] is True + assert body["input"][0]["call_id"] == "call_abc123" + + +def test_responses_call_sends_developer_items_as_system_messages() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", + input=[ # mutable-ok: the Responses API takes input as a JSON list + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Answer with exactly one word."}, + {"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]}, + ], + api_key="fw-test-key", + ) + _, _, body = _sent_request(client) + assert tuple(body["input"]) == ( + {"role": "user", "content": "Hi there"}, + {"role": "system", "content": "Answer with exactly one word.", "type": "message"}, + {"role": "user", "content": [{"type": "input_text", "text": "What is the capital of France?"}]}, + ) + + +def test_responses_call_maps_pydantic_developer_items_and_replays_pydantic_output_items() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + pydantic_input: Final = cast( + ResponseInputParam, + [ # mutable-ok: the Responses API takes input as a JSON list + EasyInputMessage(role="developer", content="Answer with exactly one word.", type="message"), + ResponseReasoningItem(id="rs_1", summary=(), type="reasoning"), + ResponseFunctionToolCall( + id="fc_1", + call_id="call_abc123", + name="get_weather", + arguments='{"city": "Paris"}', + status="completed", + type="function_call", + ), + FunctionCallOutput(type="function_call_output", call_id="call_abc123", output="21C"), + ], + ) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/accounts/fireworks/models/kimi-k3", input=pydantic_input, api_key="fw-test-key" + ) + _, _, body = _sent_request(client) + assert tuple(body["input"]) == ( + {"role": "system", "content": "Answer with exactly one word.", "type": "message"}, + {"id": "rs_1", "summary": [], "type": "reasoning"}, + { + "id": "fc_1", + "call_id": "call_abc123", + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + "type": "function_call", + }, + {"type": "function_call_output", "call_id": "call_abc123", "output": "21C"}, + ) + + +def test_file_search_tools_take_litellm_emulated_search_not_fireworks() -> None: + config: Final = FireworksAIResponsesAPIConfig() + file_search: Final = ({"type": "file_search", "vector_store_ids": ("vs_kb",)},) + function_tool: Final = ({"type": "function", "name": "get_weather", "parameters": {"type": "object"}},) + assert should_use_emulated_file_search(tools=file_search, provider_config=config) + assert not should_use_emulated_file_search(tools=function_tool, provider_config=config) + + +def test_responses_call_sends_session_affinity_for_caller_session_id() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-test-key", litellm_session_id="sess-42") + _, headers, _ = _sent_request(client) + assert headers["x-session-affinity"] == "sess-42" + + +def test_responses_call_keeps_caller_supplied_session_affinity_header() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + pinned: Final[Mapping[str, str]] = MappingProxyType({"x-session-affinity": "explicit-node"}) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input="hi", + api_key="fw-test-key", + litellm_session_id="sess-42", + extra_headers=pinned, + ) + _, headers, _ = _sent_request(client) + assert headers["x-session-affinity"] == "explicit-node" + + +def test_responses_call_maps_provider_errors_to_fireworks_ai() -> None: + client: Final = MagicMock() + request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL) + client.post.side_effect = httpx.HTTPStatusError( + "unauthorized", + request=request, + response=httpx.Response(401, text='{"error": {"message": "invalid api key"}}', request=request), + ) + with patch(HTTPX_CLIENT_FACTORY, return_value=client), pytest.raises(litellm.AuthenticationError) as raised: + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-bad-key") + assert raised.value.llm_provider == "fireworks_ai" + assert raised.value.status_code == 401 + assert "invalid api key" in str(raised.value) + + +def test_responses_call_skips_session_affinity_for_proxy_generated_session_id() -> None: + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + generated: Final[_GeneratedSessionMetadata] = {"litellm_session_id_generated": True} + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses( + model="fireworks_ai/kimi-k3", + input="hi", + api_key="fw-test-key", + litellm_session_id="generated-1", + litellm_metadata=generated, + ) + _, headers, _ = _sent_request(client) + assert "x-session-affinity" not in headers + + +@pytest.mark.parametrize( + "api_base, expected", + ( + (None, FIREWORKS_RESPONSES_URL), + ("https://api.fireworks.ai/inference/v1", FIREWORKS_RESPONSES_URL), + ("https://api.fireworks.ai/inference/v1/", FIREWORKS_RESPONSES_URL), + ("https://gateway.example.com/fireworks", "https://gateway.example.com/fireworks/responses"), + ), +) +def test_get_complete_url(api_base: str | None, expected: str) -> None: + assert FireworksAIResponsesAPIConfig().get_complete_url(api_base=api_base, litellm_params=NO_PARAMS) == expected + + +def test_responses_call_reads_fireworks_api_base_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FIREWORKS_API_BASE", "https://self-hosted.example.com/v1") + client: Final = _mock_http_client(_fireworks_response("accounts/fireworks/models/kimi-k3")) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + litellm.responses(model="fireworks_ai/kimi-k3", input="hi", api_key="fw-test-key") + url, _, _ = _sent_request(client) + assert url == "https://self-hosted.example.com/v1/responses" + + +@pytest.mark.parametrize( + "env_name", ("FIREWORKS_API_KEY", "FIREWORKS_AI_API_KEY", "FIREWORKSAI_API_KEY", "FIREWORKS_AI_TOKEN") +) +def test_validate_environment_reads_every_fireworks_key_name(monkeypatch: pytest.MonkeyPatch, env_name: str) -> None: + monkeypatch.setenv(env_name, "env-key") + headers: Final = FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=GenericLiteLLMParams() + ) + assert headers["Authorization"] == "Bearer env-key" + assert headers["Content-Type"] == "application/json" + + +def test_validate_environment_prefers_explicit_api_key(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FIREWORKS_API_KEY", "env-key") + headers: Final = FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, + model="accounts/fireworks/models/kimi-k3", + litellm_params=GenericLiteLLMParams(api_key="explicit"), + ) + assert headers["Authorization"] == "Bearer explicit" + + +def test_validate_environment_without_any_key_raises() -> None: + with pytest.raises(ValueError, match="FIREWORKS_API_KEY"): + FireworksAIResponsesAPIConfig().validate_environment( + headers=NO_HEADERS, model="accounts/fireworks/models/kimi-k3", litellm_params=None + ) + + +def test_delete_responses_maps_fireworks_message_only_body_to_deleted_result() -> None: + response_id: Final = ( + "resp_xFaIJR9Nc_OXmqKRqL78UuAGj2Te5GY5BT_knpZiMrYoNOVmu5oc2mQW1HI7hCtEYB4mcx2lEYS0DYP1U5yEQskHunuB4==" + ) + request: Final = httpx.Request("DELETE", f"{FIREWORKS_RESPONSES_URL}/{quote(response_id, safe='')}") + client: Final = MagicMock() + client.delete.return_value = httpx.Response(200, json={"message": "Response deleted successfully"}, request=request) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + result: Final = litellm.delete_responses( + response_id=response_id, custom_llm_provider="fireworks_ai", api_key="fw-test-key" + ) + assert client.delete.call_args.kwargs["url"] == str(request.url) + assert (result.id, result.object, result.deleted) == (response_id, "response", True) + + +def _fireworks_stream_response(status: str, output: tuple[Mapping[str, object], ...]) -> Mapping[str, object]: + return { + "id": "resp_htnkJ8piNKeOHkn9LfAusC38O2OgcDQs4S8trSOJ6anLeqjUDGqu2PkWmg5N", + "object": "response", + "created_at": 1788567245, + "model": "accounts/fireworks/models/kimi-k3", + "status": status, + "output": output, + "usage": None + if status == "in_progress" + else { + "input_tokens": 95, + "output_tokens": 89, + "total_tokens": 184, + "input_tokens_details": {"cached_tokens": 94}, + }, + } + + +FIREWORKS_SSE_EVENTS: Final[tuple[Mapping[str, object], ...]] = ( + {"type": "response.created", "sequence_number": 0, "response": _fireworks_stream_response("in_progress", ())}, + { + "type": "response.output_item.added", + "sequence_number": 1, + "output_index": 0, + "item": {"id": "rs_1", "type": "reasoning", "summary": []}, + }, + { + "type": "response.reasoning_summary_text.delta", + "sequence_number": 2, + "item_id": "rs_1", + "output_index": 0, + "summary_index": 0, + "delta": "pong", + }, + { + "type": "response.output_item.added", + "sequence_number": 3, + "output_index": 1, + "item": {"id": "msg_1", "type": "message", "role": "assistant", "status": "in_progress", "content": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 4, + "item_id": "msg_1", + "output_index": 1, + "content_index": 0, + "delta": "po", + }, + { + "type": "response.output_text.delta", + "sequence_number": 5, + "item_id": "msg_1", + "output_index": 1, + "content_index": 0, + "delta": "ng", + }, + { + "type": "response.completed", + "sequence_number": 6, + "response": _fireworks_stream_response( + "completed", + ( + {"id": "rs_1", "type": "reasoning", "summary": [{"type": "summary_text", "text": "pong"}]}, + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + }, + ), + ), + }, +) + + +def _sse_body(events: tuple[Mapping[str, object], ...]) -> bytes: + return b"".join(f"data: {json.dumps(dict(event))}\n\n".encode() for event in events) + b"data: [DONE]\n\n" + + +def test_streaming_responses_call_hits_native_endpoint_and_yields_every_fireworks_event() -> None: + request: Final = httpx.Request("POST", FIREWORKS_RESPONSES_URL) + client: Final = MagicMock() + client.post.return_value = httpx.Response( + 200, content=_sse_body(FIREWORKS_SSE_EVENTS), headers={"content-type": "text/event-stream"}, request=request + ) + with patch(HTTPX_CLIENT_FACTORY, return_value=client): + received: Final = tuple( + litellm.responses( + model="fireworks_ai/kimi-k3", + input="Reply with the single word pong.", + stream=True, + api_key="fw-test-key", + ) + ) + url, _, body = _sent_request(client) + assert (url, body["model"], body["stream"], client.post.call_args.kwargs["stream"]) == ( + FIREWORKS_RESPONSES_URL, + "accounts/fireworks/models/kimi-k3", + True, + True, + ) + assert tuple(event.type for event in received) == tuple(event["type"] for event in FIREWORKS_SSE_EVENTS) + assert "".join(event.delta for event in received if event.type == "response.output_text.delta") == "pong" + assert received[-1].response.usage.output_tokens == 89 diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py index 5641439aa54..41f6ad9d99d 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py @@ -56,17 +56,6 @@ def use_local_model_cost_map(): monkeypatch.undo() -@pytest.mark.parametrize("alias", KIMI_ALIASES) -def test_fireworks_kimi_raw_cost_entry_limits(use_local_model_cost_map, alias): - entry = use_local_model_cost_map.model_cost[alias] - - assert entry["litellm_provider"] == "fireworks_ai" - assert entry["max_input_tokens"] == CONTEXT_WINDOW - assert entry["max_output_tokens"] == OUTPUT_LIMIT - assert entry["max_tokens"] == OUTPUT_LIMIT - assert entry["max_output_tokens"] < entry["max_input_tokens"] - - @pytest.mark.parametrize("alias", KIMI_ALIASES) def test_fireworks_kimi_get_model_info_limits(use_local_model_cost_map, alias): model_info = use_local_model_cost_map.get_model_info(model=alias) 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/gemini/realtime/test_gemini_realtime_transformation.py b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index 3d8200bc474..2b3b6343fad 100644 --- a/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,13 +1,11 @@ import json -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock -import httpx import pytest import litellm from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig -from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents def test_gemini_realtime_transformation_session_created(): @@ -308,20 +306,6 @@ def test_gemini_realtime_transformation_generation_complete(): assert contains_audio_done_event, "Expected audio done event" -def test_gemini_3_1_flash_live_preview_model_cost_map_entry(): - for key in ( - "gemini-3.1-flash-live-preview", - "gemini/gemini-3.1-flash-live-preview", - ): - assert key in litellm.model_cost - info = litellm.model_cost[key] - assert "/v1/realtime" in info.get("supported_endpoints", []) - assert info.get("max_input_tokens") == 131072 - assert info.get("max_output_tokens") == 65536 - assert "video" in info.get("supported_modalities", []) - assert info.get("supports_function_calling") is True - - def test_gemini_realtime_tool_call_transformation(): """Test transformation of Gemini toolCall to OpenAI function_call_arguments.done format.""" config = GeminiRealtimeConfig() @@ -1845,19 +1829,6 @@ def test_is_audio_only_live_model_uses_cost_map(model, expected, patch_gemini_au assert GeminiRealtimeConfig._is_audio_only_live_model(model) == expected -def test_gemini_live_native_audio_entry_is_vertex_only(): - import json - from pathlib import Path - from typing import Final - - catalog_path: Final = Path(__file__).parents[5] / "model_prices_and_context_window.json" - catalog: Final = json.loads(catalog_path.read_text()) - vertex_key: Final = "gemini-live-2.5-flash-native-audio" - assert catalog[vertex_key]["litellm_provider"] == "vertex_ai-language-models" - assert catalog[vertex_key].get("gemini_native_audio") is True - assert "gemini/gemini-live-2.5-flash-native-audio" not in catalog, "the Gemini API does not serve this model" - - def test_is_setup_message_and_is_content_message(): config = GeminiRealtimeConfig() assert config.is_setup_message({"setup": {}}) is True diff --git a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py index cff3c6be940..4c0f5969249 100644 --- a/tests/test_litellm/llms/inception/test_inception_chat_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_chat_transformation.py @@ -231,26 +231,6 @@ def test_inception_in_provider_lists(): assert "https://api.inceptionlabs.ai/v1" in litellm.openai_compatible_endpoints -def test_inception_model_configuration(monkeypatch): - from litellm import get_model_info - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.inception_models = set() - litellm.add_known_models() - - info = get_model_info("inception/mercury-2") - assert info.get("litellm_provider") == "inception" - assert info.get("mode") == "chat" - assert info.get("max_input_tokens") == 128000 - assert info.get("input_cost_per_token") == 2.5e-07 - assert info.get("output_cost_per_token") == 7.5e-07 - assert info.get("cache_read_input_token_cost") == 2.5e-08 - assert info.get("supports_function_calling") is True - assert info.get("supports_tool_choice") is True - assert info.get("supports_response_schema") is True - - def test_inception_model_list_populated(monkeypatch): monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") litellm.model_cost = litellm.get_model_cost_map(url="") diff --git a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py index 62688a13c35..ed3f34fc744 100644 --- a/tests/test_litellm/llms/inception/test_inception_completion_transformation.py +++ b/tests/test_litellm/llms/inception/test_inception_completion_transformation.py @@ -143,24 +143,6 @@ async def test_inception_fim_async(): assert r.choices[0].text == "a + b" -def test_inception_fim_model_configuration(monkeypatch): - from litellm import get_model_info - - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - litellm.text_completion_inception_models = set() - litellm.add_known_models() - - assert ( - "text-completion-inception/mercury-edit-2" - in litellm.text_completion_inception_models - ) - info = get_model_info("text-completion-inception/mercury-edit-2") - assert info.get("litellm_provider") == "text-completion-inception" - assert info.get("mode") == "completion" - assert info.get("max_input_tokens") == 32000 - - def test_inception_fim_targets_fim_endpoint(): """ End-to-end: a FIM request must hit `/v1/fim/completions` (NOT 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/moonshot/test_moonshot_chat_transformation.py b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py index 8c8bea00dea..d484fa437ae 100644 --- a/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py +++ b/tests/test_litellm/llms/moonshot/test_moonshot_chat_transformation.py @@ -708,38 +708,6 @@ class TestKimiK26ModelRegistry: """Load directly from the bundled backup so tests don't depend on remote fetch.""" return GetModelCostMap.load_local_model_cost_map() - def test_kimi_k26_in_model_cost_map(self, model_cost_map): - """kimi-k2.6 should be present in the model cost map.""" - assert "moonshot/kimi-k2.6" in model_cost_map, "moonshot/kimi-k2.6 not found in model_cost" - - def test_kimi_k26_pricing(self, model_cost_map): - """kimi-k2.6 pricing should match official Kimi API rates.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07) - assert model_info["output_cost_per_token"] == pytest.approx(4e-06) - assert model_info["cache_read_input_token_cost"] == pytest.approx(1.6e-07) - - def test_kimi_k26_context_window(self, model_cost_map): - """kimi-k2.6 should have a 256K (262144 token) context window.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - assert model_info["max_tokens"] == 262144 - - def test_kimi_k26_capabilities(self, model_cost_map): - """kimi-k2.6 should support function calling, vision, video input, tool choice, and reasoning.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info.get("supports_function_calling") is True - assert model_info.get("supports_tool_choice") is True - assert model_info.get("supports_vision") is True - assert model_info.get("supports_video_input") is True - assert model_info.get("supports_reasoning") is True - - def test_kimi_k26_provider(self, model_cost_map): - """kimi-k2.6 should be assigned to the moonshot provider.""" - model_info = model_cost_map["moonshot/kimi-k2.6"] - assert model_info["litellm_provider"] == "moonshot" - class TestMoonshotResponseSchemaSupport: """Every model currently live on api.moonshot.ai supports json_schema @@ -762,10 +730,6 @@ class TestMoonshotResponseSchemaSupport: def model_cost_map(self): return GetModelCostMap.load_local_model_cost_map() - @pytest.mark.parametrize("model", LIVE_MODELS) - def test_live_model_supports_response_schema(self, model, model_cost_map): - assert model_cost_map[model].get("supports_response_schema") is True - def test_supports_response_schema_utility_reports_true(self, model_cost_map, monkeypatch): monkeypatch.setattr(litellm, "model_cost", model_cost_map) assert litellm.utils.supports_response_schema(model="moonshot/kimi-k2.5") is True 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/openai_like/test_cognition_provider.py b/tests/test_litellm/llms/openai_like/test_cognition_provider.py index 5c71b60e08a..d392abc6cc5 100644 --- a/tests/test_litellm/llms/openai_like/test_cognition_provider.py +++ b/tests/test_litellm/llms/openai_like/test_cognition_provider.py @@ -110,22 +110,6 @@ class TestCognitionProviderIdentity: class TestCognitionCostTracking: - @pytest.mark.parametrize( - "model, input_cost, output_cost, cache_read_cost", - [ - ("cognition/swe-1.6", 5e-07, 2.5e-06, 2e-07), - ("cognition/swe-1.7", 5e-07, 2.5e-06, 2e-07), - ("cognition/swe-1.7-lightning", 2.5e-06, 1.25e-05, 1e-06), - ], - ) - def test_cost_map_entries(self, model: str, input_cost: float, output_cost: float, cache_read_cost: float): - info = litellm.get_model_info(model=model) - - assert info["litellm_provider"] == "cognition" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cache_read_cost @pytest.mark.parametrize( "model, expected_prompt_cost, expected_completion_cost", diff --git a/tests/test_litellm/llms/openai_like/test_json_providers.py b/tests/test_litellm/llms/openai_like/test_json_providers.py index c8743e1809d..d84cc8d3237 100644 --- a/tests/test_litellm/llms/openai_like/test_json_providers.py +++ b/tests/test_litellm/llms/openai_like/test_json_providers.py @@ -2,10 +2,9 @@ Tests for JSON-based provider configuration system. """ -import json import os import sys -from unittest.mock import MagicMock, patch +from unittest.mock import patch try: import pytest @@ -318,25 +317,6 @@ class TestDarkbloom: assert config is not None assert config.custom_llm_provider == "darkbloom" - def test_darkbloom_model_cost_map(self): - with open( - os.path.join(workspace_path, "model_prices_and_context_window.json") - ) as f: - model_cost = json.load(f) - - expected_models = { - "darkbloom/gemma-4-26b": (3e-08, 1.65e-07), - "darkbloom/gpt-oss-20b": (1.45e-08, 7e-08), - } - for model, (input_cost, output_cost) in expected_models.items(): - assert model in model_cost - assert model_cost[model]["litellm_provider"] == "darkbloom" - assert model_cost[model]["max_output_tokens"] == 32768 - assert model_cost[model]["supports_function_calling"] is True - assert model_cost[model]["supports_tool_choice"] is True - assert model_cost[model]["input_cost_per_token"] == input_cost - assert model_cost[model]["output_cost_per_token"] == output_cost - class TestPublicAIIntegration: """Integration tests for PublicAI provider""" diff --git a/tests/test_litellm/llms/openai_like/test_libertai_provider.py b/tests/test_litellm/llms/openai_like/test_libertai_provider.py index fdbe3046e9b..c17eaf7c87f 100644 --- a/tests/test_litellm/llms/openai_like/test_libertai_provider.py +++ b/tests/test_litellm/llms/openai_like/test_libertai_provider.py @@ -59,23 +59,6 @@ class TestLibertAIProviderConfig: assert api_base == "https://custom.example.com/v1" assert api_key == "sk-test" - def test_libertai_model_cost_map(self): - """Test that libertai models are present in the model cost map""" - model_cost = litellm.model_cost - - assert "libertai/qwen3.6-27b" in model_cost - info = model_cost["libertai/qwen3.6-27b"] - assert info["litellm_provider"] == "libertai" - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - - # thinking variants are marked as reasoning models - assert ( - model_cost["libertai/qwen3.6-27b-thinking"].get("supports_reasoning") - is True - ) - def test_libertai_router_config(self): """Test that libertai can be used in Router configuration""" from litellm import Router @@ -95,20 +78,6 @@ class TestLibertAIProviderConfig: assert len(router.model_list) == 1 assert router.model_list[0]["model_name"] == "libertai-chat" - def test_libertai_model_modes(self): - """Chat models carry mode 'chat'; the embedding model carries mode 'embedding'.""" - model_cost = litellm.model_cost - - # chat model - assert model_cost["libertai/qwen3.6-27b"]["mode"] == "chat" - - # embedding model (bge-m3) must be normalized to mode 'embedding' so - # /embeddings routing and the supported-endpoints matrix stay consistent - assert "libertai/bge-m3" in model_cost - bge = model_cost["libertai/bge-m3"] - assert bge["litellm_provider"] == "libertai" - assert bge["mode"] == "embedding" - def test_libertai_supported_endpoints_matrix(self): """The runtime-served backup matrix (GET /public/supported_endpoints) lists libertai.""" import json diff --git a/tests/test_litellm/llms/openai_like/test_meta_provider.py b/tests/test_litellm/llms/openai_like/test_meta_provider.py index 11b78828da6..c79e4b77cc5 100644 --- a/tests/test_litellm/llms/openai_like/test_meta_provider.py +++ b/tests/test_litellm/llms/openai_like/test_meta_provider.py @@ -193,19 +193,6 @@ class TestMetaAnthropicMessages: class TestMuseSparkModelInfo: - def test_muse_spark_pricing_and_capabilities(self): - info = litellm.get_model_info("meta/muse-spark-1.1") - - assert info["litellm_provider"] == "meta" - assert info["input_cost_per_token"] == 1.25e-06 - assert info["output_cost_per_token"] == 4.25e-06 - assert info["cache_read_input_token_cost"] == 1.5e-07 - assert info["max_input_tokens"] == 1048576 - assert info["supports_reasoning"] is True - assert info["supports_web_search"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True def test_muse_spark_cost_calculation(self): from litellm import completion_cost diff --git a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py index 6a6271e95e2..6ca7072e7ab 100644 --- a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py +++ b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py @@ -3,7 +3,6 @@ Unit tests for Perplexity embedding transformation logic. """ import base64 -import json import struct from unittest.mock import MagicMock @@ -298,25 +297,3 @@ class TestPerplexityEmbeddingProviderConfig: ) assert config is not None assert isinstance(config, PerplexityEmbeddingConfig) - - -class TestPerplexityEmbeddingModelInfo: - """Test that Perplexity embedding models are in model_prices_and_context_window.""" - - def test_model_info_available(self): - import litellm - - info = litellm.get_model_info("perplexity/pplx-embed-v1-0.6b") - assert info is not None - assert info["mode"] == "embedding" - assert info["max_input_tokens"] == 32768 - assert info["output_vector_size"] == 1024 - - def test_model_info_4b_available(self): - import litellm - - info = litellm.get_model_info("perplexity/pplx-embed-v1-4b") - assert info is not None - assert info["mode"] == "embedding" - assert info["max_input_tokens"] == 32768 - assert info["output_vector_size"] == 2560 diff --git a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py index 6630039e92e..7556b215e66 100644 --- a/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/test_litellm/llms/perplexity/test_perplexity_cost_calculator.py @@ -26,7 +26,6 @@ from litellm.types.utils import ( Usage, PromptTokensDetailsWrapper, ) -from litellm.utils import get_model_info class TestPerplexityCostCalculator: @@ -317,21 +316,6 @@ class TestPerplexityCostCalculator: assert math.isclose(total_cost, expected_total, rel_tol=1e-6) - def test_model_info_access(self): - """Test that model info correctly returns the new cost fields.""" - model_info = get_model_info( - model="sonar-deep-research", custom_llm_provider="perplexity" - ) - - # Check that the new fields are accessible - assert "citation_cost_per_token" in model_info - assert model_info["citation_cost_per_token"] == 2e-6 - assert model_info["search_context_cost_per_query"] == { - "search_context_size_low": 0.005, - "search_context_size_medium": 0.005, - "search_context_size_high": 0.005, - } - @pytest.mark.parametrize("citation_tokens", [0, 10, 25, 100]) @pytest.mark.parametrize("search_queries", [0, 1, 5, 10]) @pytest.mark.parametrize("reasoning_tokens", [0, 15, 30]) @@ -477,37 +461,6 @@ class TestPerplexityCostCalculator: assert math.isclose(prompt_cost, expected_prompt, rel_tol=1e-9) assert math.isclose(completion_cost, expected_completion, rel_tol=1e-9) - @pytest.mark.parametrize( - "model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read", - [ - ("deepseek-v4-flash-0731", 0.13, 0.26, 0.028), - ("glm-5.2", 1.4, 4.4, 0.14), - ("kimi-k3", 3.0, 15.0, 0.3), - ("kimi-k2.7-code", 0.95, 4.0, 0.19), - ], - ) - def test_agent_api_entries_carry_perplexity_published_rates( - self, model_id, usd_per_1m_input, usd_per_1m_output, usd_per_1m_cache_read - ): - """The Agent API third-party models are priced from Perplexity's own catalog - (GET https://api.perplexity.ai/v1/models, `pricing` in usd_per_1m_tokens). - Perplexity's model id already starts with `perplexity/`, so the cost-map key - doubles the prefix. Regression: glm-5.2 shipped glm-5.3's 0.26 cache-read rate, - copied from the neighbouring catalog row, an 86% overcharge on cached input. - """ - info = get_model_info( - model=f"perplexity/{model_id}", custom_llm_provider="perplexity" - ) - - assert info["key"] == f"perplexity/perplexity/{model_id}" - assert info["litellm_provider"] == "perplexity" - assert info["mode"] == "responses" - assert math.isclose(info["input_cost_per_token"], usd_per_1m_input / 1e6, rel_tol=1e-9) - assert math.isclose(info["output_cost_per_token"], usd_per_1m_output / 1e6, rel_tol=1e-9) - assert math.isclose( - info["cache_read_input_token_cost"], usd_per_1m_cache_read / 1e6, rel_tol=1e-9 - ) - def test_agent_api_fallback_rates_price_a_response_without_metered_cost(self): """Perplexity meters cost on the response, but when `usage.cost` is absent the calculator falls back to the mapped per-token rates. Regression: that fallback diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 38fde3caa63..24df3214da3 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -35,7 +35,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest - from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) @@ -178,6 +177,197 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): sync_client.post.assert_not_called() +def test_create_batch_sync_does_not_resolve_publisher_models(): + """Publisher-model jobs must not incur the endpoint-resolution GET, and the job model must + stay the publisher path untouched.""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get") as safe_get, + ): + out = h.create_batch( + _is_async=False, + create_batch_data=CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert isinstance(out, LiteLLMBatch) + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001" + safe_get.assert_not_called() + + +ENDPOINT_ID = "7768560373388541952" +ENDPOINT_CREATE_DATA = { + "input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid" +} +TUNED_MODEL_RESOURCE = f"projects/{PROJECT}/locations/{LOCATION}/models/1234509876" + + +def _endpoint_get_response(deployed_models: list | None = None) -> MagicMock: + resp = MagicMock() + resp.status_code = 200 + resp.json.return_value = { + "name": f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + "deployedModels": ( + deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}] + ), + } + return resp + + +def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): + """A fine-tuned Gemini file id must produce a batch job against the endpoint's deployed + tuned model resource; the v1 batch API rejects endpoint resources in `model` (LIT-6899).""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()) as safe_get, + ): + out = h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert isinstance(out, LiteLLMBatch) + get_args, get_kwargs = safe_get.call_args + assert get_args[1] == ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}" + ) + assert get_kwargs["headers"]["Authorization"] == f"Bearer {TOKEN}" + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == TUNED_MODEL_RESOURCE + + +@pytest.mark.parametrize( + "api_base, expected", + [ + ( + None, + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal", + f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal/v1", + f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal/vertex", + f"https://proxy.internal/vertex/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ], +) +def test_build_endpoint_resolution_url(api_base, expected): + """A custom api_base must replace the Google host for the endpoint-resolution GET without + producing a malformed url (no ':' grafting, no doubled /v1).""" + url = VertexAIBatchPrediction._build_endpoint_resolution_url( + api_base=api_base, + model=f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + vertex_location=LOCATION, + ) + assert url == expected + + +def test_create_batch_sync_endpoint_resolution_error_raises(): + h = _make_handler() + client = MagicMock() + resolve_response = MagicMock() + resolve_response.status_code = 404 + resolve_response.text = "endpoint not found" + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=resolve_response), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 404 + client.post.assert_not_called() + + +def test_create_batch_custom_endpoint_raises_400_without_io(): + """custom_endpoint deployments have no Vertex batch surface; creating a job would target a + nonexistent publisher model, so the handler must 400 before any auth or HTTP work (LIT-6899).""" + h = _make_handler() + client = MagicMock() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + assert exc_info.value.status_code == 400 + assert "custom_endpoint" in str(exc_info.value) + h._ensure_access_token.assert_not_called() + client.post.assert_not_called() + + +def test_create_batch_sync_endpoint_without_deployed_model_raises_400(): + h = _make_handler() + client = MagicMock() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response(deployed_models=[])), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + assert "no deployed model" in str(exc_info.value) + client.post.assert_not_called() + + def test_create_batch_sync_httpstatuserror_propagates(): """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the sync create path must surface that error, not swallow it.""" diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index ccb2d7e310d..232c6413e78 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -34,6 +34,12 @@ INPUT_FILE = ( "models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8" ) +ENDPOINT_ID = "7768560373388541952" +ENDPOINT_INPUT_FILE = ( + f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/" + "e9412502-2c91-42a6-8e61-f5c294cc0fc8" +) + # =========================================================================== # # transform_openai_batch_request_to_vertex_ai_batch_request @@ -67,6 +73,41 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +def test_transform_openai_request_fine_tuned_endpoint_builds_endpoint_resource(): + """A fine-tuned Gemini file id (endpoints/) must target the endpoint resource, + not a nonexistent publisher model (LIT-6899).""" + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + vertex_location="us-central1", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}" + + +def test_transform_openai_request_fine_tuned_endpoint_defaults_location(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}" + + +def test_transform_openai_request_fine_tuned_endpoint_without_project_raises_400(): + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": ENDPOINT_INPUT_FILE}) + assert exc_info.value.status_code == 400 + assert "vertex_project" in str(exc_info.value) + + +def test_transform_openai_request_publisher_model_ignores_project_and_location(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": INPUT_FILE}, + vertex_project="my-project", + vertex_location="europe-west4", + ) + assert job["model"] == "publishers/google/models/gemini-1.5-flash-001" + + @pytest.mark.parametrize( "input_file_id", [ @@ -321,6 +362,35 @@ def test_get_model_from_gcs_file_no_publishers_raises_400(): assert exc_info.value.status_code == 400 +def test_get_model_from_gcs_file_fine_tuned_endpoint(): + """The whole endpoint id must survive parsing; the old 3-segment publishers/ parse dropped it.""" + assert T._get_model_from_gcs_file(ENDPOINT_INPUT_FILE) == f"endpoints/{ENDPOINT_ID}" + + +def test_get_model_from_gcs_file_publisher_path_wins_over_endpoints_prefix(): + """A bucket prefix containing endpoints/ must not override the publisher model path + LiteLLM appended after it.""" + uri = "gs://bucket/team-endpoints/999/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/uuid" + assert T._get_model_from_gcs_file(uri) == "publishers/google/models/gemini-1.5-flash-001" + + +def test_get_model_from_gcs_file_last_endpoints_segment_wins(): + """With no publisher path, the endpoint id closest to the file (last occurrence) is the one + LiteLLM stored; an earlier prefix segment must not shadow it.""" + uri = f"gs://bucket/endpoints/999/litellm-vertex-files/endpoints/{ENDPOINT_ID}/uuid" + assert T._get_model_from_gcs_file(uri) == f"endpoints/{ENDPOINT_ID}" + + +def test_get_model_from_gcs_file_non_numeric_endpoints_segment_raises_400(): + with pytest.raises(VertexAIError) as exc_info: + T._get_model_from_gcs_file("gs://bucket/endpoints/not-a-number/file-uuid") + assert exc_info.value.status_code == 400 + + +def test_get_bare_model_name_from_gcs_file_fine_tuned_endpoint(): + assert T.get_bare_model_name_from_gcs_file(ENDPOINT_INPUT_FILE) == ENDPOINT_ID + + # =========================================================================== # # is_unmanaged_gcs_batch_input_file_id # =========================================================================== # @@ -334,6 +404,8 @@ def test_get_model_from_gcs_file_no_publishers_raises_400(): ("file-abc123", False), ("gs://bucket/no-model-here.jsonl", False), ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + (ENDPOINT_INPUT_FILE, True), + ("gs://bucket/endpoints/not-a-number/file-uuid", False), ], ) def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 3c2d56997b7..8a249820cbd 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -159,6 +159,91 @@ class TestCreateFileUrl: assert "?" not in object_name +class TestBatchObjectNaming: + def test_should_store_publisher_model_under_publishers_path(self, config): + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini-2.5-flash"}}]) + assert object_name.startswith("litellm-vertex-files/publishers/google/models/gemini-2.5-flash/") + + def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config): + """A numeric endpoint id must not be filed under publishers/google/models/gemini/, + which the batch transformation later mangles into a nonexistent publisher model (LIT-6899).""" + object_name = config._get_gcs_object_name_from_batch_jsonl( + [{"body": {"model": "gemini/7768560373388541952"}}] + ) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "publishers" not in object_name + + def test_should_store_bare_numeric_endpoint_under_endpoints_path(self, config): + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "7768560373388541952"}}]) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + + def test_deployment_model_overrides_jsonl_body_model(self, config): + """The stored path decides which Vertex model the batch later runs against with the + deployment's credentials, so a user-crafted JSONL body.model must not be able to redirect + an authorized deployment to a different endpoint.""" + object_name = config._get_gcs_object_name_from_batch_jsonl( + [{"body": {"model": "9999999999999999999"}}], + deployment_model="vertex_ai/gemini/7768560373388541952", + ) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "9999999999999999999" not in object_name + + def test_url_derives_object_path_from_configured_model(self, config): + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={ + "gcs_bucket_name": "my-bucket", + "model": "vertex_ai/gemini/7768560373388541952", + }, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "9999999999999999999"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "9999999999999999999" not in object_name + + +class TestCustomEndpointBatchUpload: + def test_should_reject_batch_upload_for_custom_endpoint_deployment(self, config): + """custom_endpoint deployments have no Vertex batch surface; the upload must 400 instead + of staging a file that can only produce a doomed batch job (LIT-6899).""" + from litellm.llms.vertex_ai.common_utils import VertexAIError + + with pytest.raises(VertexAIError) as exc_info: + config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "openai/gemma-2-2b-it"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + assert exc_info.value.status_code == 400 + assert "custom_endpoint" in str(exc_info.value) + + def test_should_allow_non_batch_upload_for_custom_endpoint_deployment(self, config): + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + data={ + "file": ("notes.txt", b"hello", "text/plain"), + "purpose": "assistants", + }, + ) + assert "/b/my-bucket/" in url + + class TestTransformRetrieveFile: def test_should_build_correct_gcs_metadata_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 6ba8706b0d8..04e46eab1b7 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -136,19 +136,6 @@ class TestVertexAIVideoConfig: # Should NOT include endpoint assert not url.endswith(":predictLongRunning") - def test_veo_31_lite_model_cost_entries_match_pricing(self): - for path in (ROOT_MODEL_COST_PATH, BACKUP_MODEL_COST_PATH): - model_cost = _load_model_cost_map(path) - info = model_cost.get(VEO_31_LITE_VERTEX_MODEL) - - assert info is not None, f"{VEO_31_LITE_VERTEX_MODEL} missing from {path}" - assert info["litellm_provider"] == "vertex_ai-video-models" - assert info["mode"] == "video_generation" - assert info["max_input_tokens"] == 1024 - assert info["output_cost_per_second"] == 0.05 - assert info["output_cost_per_second_1080p"] == 0.08 - assert info["supported_modalities"] == ["text", "image"] - def test_veo_31_lite_provider_routing_from_local_model_map( self, monkeypatch: pytest.MonkeyPatch ): diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py index 83c3bf1ecef..e591c1ae682 100644 --- a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py +++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py @@ -63,11 +63,6 @@ TIER_COST_FIELDS = ( "output_cost_per_token_above_200k_tokens", "cache_read_input_token_cost_above_200k_tokens", ) -STALE_TIER_FIELDS = ( - "input_cost_per_token_above_128k_tokens", - "output_cost_per_token_above_128k_tokens", - "cache_read_input_token_cost_above_128k_tokens", -) def expected_retirement_date(slug: str) -> str: @@ -102,11 +97,10 @@ def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str): assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug) -@pytest.mark.parametrize("slug", REDIRECTED_SLUGS) -def test_no_slug_keeps_the_superseded_128k_tier(cost_map: dict, slug: str): - """The 128k tier belonged to the retired model; grok-4.3 tiers at 200k.""" - for field in STALE_TIER_FIELDS: - assert field not in cost_map[slug], field +def test_a_live_xai_model_is_untouched(cost_map: dict): + """Guard against the repricing leaking onto models xAI still serves directly.""" + assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"] + assert "deprecation_date" not in cost_map["xai/grok-4.6"] @pytest.mark.parametrize("slug", REDIRECTED_SLUGS) @@ -116,12 +110,7 @@ def test_redirected_slug_carries_the_target_tier_rates(cost_map: dict, slug: str entry = cost_map[slug] for field in TIER_COST_FIELDS: assert entry[field] == target[field], field - - -def test_a_live_xai_model_is_untouched(cost_map: dict): - """Guard against the repricing leaking onto models xAI still serves directly.""" - assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"] - assert "deprecation_date" not in cost_map["xai/grok-4.6"] + assert {k for k in entry if "_above_" in k} == {k for k in target if "_above_" in k} def test_both_cost_maps_agree_on_the_redirected_slugs(): diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index 38ddac8d510..069ac5727f6 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -2,11 +2,9 @@ Tests for Z.AI (Zhipu AI) provider - GLM models """ -import json import math import pytest -import respx import litellm from litellm import completion @@ -57,32 +55,9 @@ def test_zai_in_provider_lists(): assert "zai" in litellm.provider_list -def test_zai_models_in_model_cost(local_model_cost_map): - """Test that ZAI models are in the model cost map""" - - zai_models = [ - "zai/glm-4.7", - "zai/glm-4.6", - "zai/glm-4.5", - "zai/glm-4.5v", - "zai/glm-4.5-x", - "zai/glm-4.5-air", - "zai/glm-4.5-airx", - "zai/glm-4-32b-0414-128k", - "zai/glm-4.5-flash", - ] - - for model in zai_models: - assert model in litellm.model_cost, f"Model {model} not found in model_cost" - assert litellm.model_cost[model]["litellm_provider"] == "zai" - - def test_zai_glm46_cost_calculation(local_model_cost_map): """Test the cost calculation for glm-4.6""" - key = "zai/glm-4.6" - info = litellm.model_cost[key] - prompt_cost, completion_cost = cost_per_token( model="zai/glm-4.6", prompt_tokens=1000000, # 1M tokens @@ -94,26 +69,6 @@ def test_zai_glm46_cost_calculation(local_model_cost_map): assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) -def test_zai_flash_model_is_free(local_model_cost_map): - """Test that glm-4.5-flash has zero cost""" - - key = "zai/glm-4.5-flash" - info = litellm.model_cost[key] - - assert info["input_cost_per_token"] == 0 - assert info["output_cost_per_token"] == 0 - - -def test_glm47_supports_reasoning(local_model_cost_map): - """Test that GLM-4.7 supports reasoning""" - - key = "zai/glm-4.7" - assert key in litellm.model_cost, f"Model {key} not found in model_cost" - - info = litellm.model_cost[key] - assert info["supports_reasoning"] is True - - def test_glm47_cost_calculation(local_model_cost_map): """Test cost calculation for GLM-4.7""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 44eb7795659..20f4719e4bf 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -11,6 +11,7 @@ from starlette.datastructures import Headers from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, + UnloadableEntitlementError, _is_mcp_admitted_user_subject, ) from litellm.proxy._types import ( @@ -4396,6 +4397,147 @@ class TestAgentMCPPermissions: ) assert sorted(result) == ["tool_a", "tool_b"] + def _agent_object_permission(self, *, toolset_ids, servers=(), tool_permissions=None): + agent_object_permission = MagicMock() + agent_object_permission.mcp_servers = list(servers) + agent_object_permission.mcp_access_groups = [] + agent_object_permission.mcp_tool_permissions = tool_permissions + agent_object_permission.mcp_toolsets = list(toolset_ids) + return agent_object_permission + + def _mock_manager_with_toolsets(self, toolset_perms): + mock_manager = MagicMock() + mock_manager.expand_permission_list = MagicMock(side_effect=lambda servers: list(servers)) + mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {}) + mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value=toolset_perms) + return mock_manager + + def _agent_toolset_patches(self, agent_object_permission, mock_manager): + return ( + patch.object( # test-quality-ok: stub the agent perm loader; the resolver reads module globals with no injection seam + MCPRequestHandler, "_get_agent_object_permission", AsyncMock(return_value=agent_object_permission) + ), + patch( # test-quality-ok: isolate the MCP registry, same seam as the sibling toolset tests + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager", + mock_manager, + ), + patch.object( # test-quality-ok: access-group lookup hits the DB, not under test here + MCPRequestHandler, "_get_mcp_servers_from_access_groups", AsyncMock(return_value=[]) + ), + ) + + async def test_get_allowed_mcp_servers_for_agent_includes_toolset_servers(self): + """An agent granted only mcp_toolsets reaches the toolset's servers, exactly as a + key, team, or org granted only toolsets does""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"], servers=["server-direct"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + result = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) + + assert sorted(result) == ["server-a", "server-direct"] + mock_manager.resolve_toolset_tool_permissions.assert_awaited_once_with(toolset_ids=["toolset-1"]) + + async def test_get_allowed_mcp_servers_toolset_only_agent_caps_key_servers(self): + """Regression: an agent whose only grant is a toolset used to resolve to [] and place + no ceiling at all, so a key bound to it kept every server the key itself granted""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["lookup_status"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + stack.enter_context( + patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[]) + ) + ) + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == ["server-a"] + + async def test_get_allowed_mcp_servers_agent_dangling_toolset_denies(self): + """An agent toolset that resolves to nothing is a known restriction with unknown + contents: deny, never fall through to the key's own servers""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-gone"]) + mock_manager = self._mock_manager_with_toolsets({}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + with pytest.raises(UnloadableEntitlementError): + await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth) + stack.enter_context( + patch.object( # test-quality-ok: key resolution has its own tests; pin its grants here + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server-a", "server-b"]) + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it empty here + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", AsyncMock(return_value=[]) + ) + ) + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + assert result == [] + + async def test_get_agent_tool_permissions_for_server_unions_direct_and_toolset_tools(self): + """The agent's tool ceiling on a server is its direct tool grants plus the tools its + toolsets grant there, and None only when neither names the server""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission( + toolset_ids=["toolset-1"], tool_permissions={"server-a": ["tool_direct"]} + ) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_via_toolset"], "server-b": ["tool_b"]}) + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + server_a_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-a", user_api_key_auth) + server_b_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-b", user_api_key_auth) + server_c_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server("server-c", user_api_key_auth) + + assert sorted(server_a_tools) == ["tool_direct", "tool_via_toolset"] + assert server_b_tools == ["tool_b"] + assert server_c_tools is None + + async def test_get_allowed_tools_for_server_toolset_only_agent_caps_key_tools(self): + """Regression: a key allowing [tool_a, tool_b] bound to an agent whose toolset grants + only tool_a on the server ends with [tool_a]; the toolset used to be ignored""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", agent_id="agent-toolsets") + agent_object_permission = self._agent_object_permission(toolset_ids=["toolset-1"]) + mock_manager = self._mock_manager_with_toolsets({"server-a": ["tool_a"]}) + key_perm = MagicMock() + key_perm.mcp_tool_permissions = {"server-a": ["tool_a", "tool_b"]} + key_perm.mcp_toolsets = [] + + with contextlib.ExitStack() as stack: + for patcher in self._agent_toolset_patches(agent_object_permission, mock_manager): + stack.enter_context(patcher) + stack.enter_context( + patch.object( # test-quality-ok: stub the key perm loader; the resolver reads module globals with no injection seam + MCPRequestHandler, "_get_key_object_permission", return_value=key_perm + ) + ) + stack.enter_context( + patch.object( # test-quality-ok: team resolution has its own tests; pin it absent here + MCPRequestHandler, "_get_team_object_permission", AsyncMock(return_value=None) + ) + ) + result = await MCPRequestHandler.get_allowed_tools_for_server("server-a", user_api_key_auth) + + assert result == ["tool_a"] + async def test_get_agent_object_permission_uses_shared_helper(self): """``_get_agent_object_permission`` must resolve the agent's ``object_permission_id`` and then defer to the shared diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index e20f6646310..355b3bfd30e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -12,28 +12,39 @@ import base64 import json from datetime import datetime, timedelta, timezone from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest +from prisma.models import LiteLLM_MCPServerTable as PrismaMCPServer from litellm.proxy._experimental.mcp_server.db import ( _decode_user_credential, _prepare_mcp_server_data, + create_mcp_server, decrypt_credentials, encrypt_credentials, + get_all_mcp_servers, + get_mcp_servers, + get_mcp_submissions, get_user_credential, get_user_oauth_credential, is_oauth_credential_expired, list_user_oauth_credentials, resolve_valid_user_oauth_token, + rotate_mcp_server_credentials_master_key, rotate_mcp_user_credentials_master_key, rotate_mcp_user_env_vars_master_key, store_user_credential, store_user_oauth_credential, + update_mcp_server, ) -from litellm.proxy._types import NewMCPServerRequest, UpdateMCPServerRequest +from litellm.proxy._types import LiteLLM_MCPServerTable, NewMCPServerRequest, UpdateMCPServerRequest from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + SecretMapDecodeError, + decode_secret_map, decrypt_value_helper, + encrypt_secret_map, encrypt_value_helper, ) from litellm.types.mcp import MCPAuth, MCPTransport @@ -44,6 +55,7 @@ SALT_KEY = "test-salt-key-for-byok-credential-tests-1234" @pytest.fixture(autouse=True) def _set_salt_key(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": "xsalsa20-poly1305"}) def _make_prisma_with_existing(row): @@ -368,6 +380,169 @@ def test_client_private_key_encrypted_at_rest(): assert decrypted["client_secret"] == "shh" +@pytest.fixture(params=["xsalsa20-poly1305", "aes-256-gcm"]) +def map_algorithm(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> str: + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": request.param}) + return request.param + + +def _prisma_map_row(data: dict[str, object], quoted: bool = False) -> PrismaMCPServer: + return PrismaMCPServer.model_validate({ + "transport": "http", "mcp_access_groups": [], "allowed_tools": [], "extra_headers": [], "args": [], + "allow_all_keys": False, "available_on_public_internet": True, "delegate_auth_to_upstream": False, + "oauth_passthrough": False, "per_server_oauth_discovery": False, "is_byok": False, "byok_description": [], + **data, + **{field: json.dumps(data[field]) for field in ("static_headers", "env") if quoted and data.get(field)}, + }) + + +class _MapTable: + def __init__(self, *rows: dict[str, object], quoted: bool = False) -> None: + self.rows = {row["server_id"]: row for row in rows} + self.quoted = quoted + + async def create(self, *, data: dict[str, object]) -> PrismaMCPServer: + self.rows = {**self.rows, data["server_id"]: dict(data)} + return _prisma_map_row(data, self.quoted) + + async def update(self, *, where: dict[str, str], data: dict[str, object]) -> PrismaMCPServer: + return await self.create(data={**self.rows[where["server_id"]], **data}) + + async def find_many(self, where: object = None) -> list[PrismaMCPServer]: + return [_prisma_map_row(row, self.quoted) for row in self.rows.values()] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["static_headers", "env"]) +@pytest.mark.parametrize("quoted", [False, True]) +async def test_secret_maps_create_update_round_trip(map_algorithm: str, field: str, quoted: bool) -> None: + table: Final = _MapTable(quoted=quoted) + prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + original: Final = {"TOKEN": " sensitive-secret\n", "PREFIX": "v2:gcm:literal", "TEMPLATE": "Bearer ${TOKEN}"} + create: Final = NewMCPServerRequest.model_validate({ + "server_id": "srv-map", "transport": "http", "url": "https://up.example.com/mcp", field: original, + }) + created: Final = await create_mcp_server(prisma, create, touched_by="test") + first: Final = table.rows["srv-map"][field] + assert isinstance(first, str) and isinstance(json.loads(first), str) + assert json.loads(first).startswith("v2:gcm:") is (map_algorithm == "aes-256-gcm") + assert "sensitive-secret" not in first and "TEMPLATE" not in first + assert getattr(created, field) == original == getattr(create, field) + assert decode_secret_map(first, key=field) == original + replacement: Final = {**original, "TOKEN": "updated-sensitive-secret"} + update: Final = UpdateMCPServerRequest.model_validate({"server_id": "srv-map", field: replacement}) + updated: Final = await update_mcp_server(prisma, update, touched_by="test") + second: Final = table.rows["srv-map"][field] + assert second != first and "updated-sensitive-secret" not in second + assert decode_secret_map(second, key=field) == replacement + assert getattr(updated, field) == replacement == getattr(update, field) + assert original["TOKEN"] == " sensitive-secret\n" + omitted: Final = await update_mcp_server(prisma, UpdateMCPServerRequest(server_id="srv-map"), touched_by="test") + assert table.rows["srv-map"][field] == second and getattr(omitted, field) == replacement + cleared: Final = await update_mcp_server( + prisma, UpdateMCPServerRequest.model_validate({"server_id": "srv-map", field: {}}), touched_by="test" + ) + assert table.rows["srv-map"][field] == "{}" and getattr(cleared, field) == {} + + +@pytest.mark.parametrize("field", ["static_headers", "env"]) +@pytest.mark.parametrize("as_json", [False, True]) +def test_secret_map_legacy_model_read_preserves_exact_values(field: str, as_json: bool) -> None: + original: Final = {"PREFIX": "v2:gcm:literal", "SPACE": " secret\n", "TEMPLATE": "${TOKEN}", "B64": "YWJjZA=="} + incoming: Final = { + "server_id": "srv-map", "transport": "http", field: json.dumps(original) if as_json else original, + } + snapshot: Final = json.dumps(incoming) + parsed: Final = LiteLLM_MCPServerTable.model_validate(incoming) + assert getattr(parsed, field) == original + assert json.dumps(incoming) == snapshot + assert LiteLLM_MCPServerTable.model_validate(parsed.model_dump()).model_dump() == parsed.model_dump() + empty: Final = LiteLLM_MCPServerTable.model_validate({"server_id": "srv-map", "transport": "http", field: None}) + assert getattr(empty, field) == ({} if field == "env" else None) + + +@pytest.mark.parametrize("field", ["static_headers", "env"]) +@pytest.mark.parametrize("failure", ["wrong-key", "corrupt", "invalid-values", "invalid-shape", "invalid-json"]) +def test_secret_map_model_read_fails_closed(map_algorithm: str, field: str, failure: str) -> None: + plaintext: Final = {"invalid-values": '{"TOKEN": ["sensitive-secret"]}', "invalid-shape": '["sensitive-secret"]', + "invalid-json": "sensitive-secret"}.get(failure, '{"TOKEN": "sensitive-secret"}') + ciphertext: Final = encrypt_value_helper( + plaintext, new_encryption_key="wrong-map-key" if failure == "wrong-key" else None + ) + stored: Final = json.dumps(ciphertext[:-8] if failure == "corrupt" else ciphertext) + with pytest.raises(SecretMapDecodeError) as exc: + LiteLLM_MCPServerTable.model_validate({"server_id": "srv-map", "transport": "http", field: stored}) + assert field in str(exc.value) and "LITELLM_SALT_KEY" in str(exc.value) + assert all(secret not in str(exc.value) for secret in (plaintext, ciphertext, "sensitive-secret")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field,other", [("static_headers", "env"), ("env", "static_headers")]) +async def test_secret_map_rotation_migrates_rekeys_and_preserves_corrupt( + map_algorithm: str, field: str, other: str, monkeypatch: pytest.MonkeyPatch +) -> None: + values: Final = {"TOKEN": "rotation-sensitive-secret", "TEMPLATE": "Bearer ${TOKEN}"} + old: Final = encrypt_secret_map(values) + corrupt: Final = json.dumps(json.loads(old)[:-8]) + table: Final = _MapTable( + {"server_id": "broken", field: corrupt, other: old}, + {"server_id": "legacy", field: json.dumps(values), other: "{}"}, + {"server_id": "encrypted", field: old, other: None}, + ) + prisma: Final = SimpleNamespace(db=SimpleNamespace( + litellm_mcpservertable=table, litellm_mcpserveroauthclient=SimpleNamespace(find_many=AsyncMock(return_value=[])) + )) + await rotate_mcp_server_credentials_master_key(prisma, touched_by="test", new_master_key="rotated-map-key") + assert table.rows["broken"][field] == corrupt + assert table.rows["legacy"][other] == "{}" and table.rows["encrypted"][other] is None + for server_id, map_field in (("broken", other), ("legacy", field), ("encrypted", field)): + stored: Final = table.rows[server_id][map_field] + assert isinstance(json.loads(stored), str) and stored != old and "rotation-sensitive-secret" not in stored + with pytest.raises(SecretMapDecodeError): + decode_secret_map(stored, key=map_field) + monkeypatch.setenv("LITELLM_SALT_KEY", "rotated-map-key") + for server_id, map_field in (("broken", other), ("legacy", field), ("encrypted", field)): + assert decode_secret_map(table.rows[server_id][map_field], key=map_field) == values + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reader", [get_all_mcp_servers, get_mcp_servers, get_mcp_submissions]) +@pytest.mark.parametrize("field", ["static_headers", "env"]) +async def test_bulk_reads_isolate_corrupt_secret_maps(reader, field, map_algorithm, caplog): + secret = {"TOKEN": "bulk-sensitive-secret"} + encrypted = encrypt_secret_map(secret) + corrupt = encrypt_secret_map(secret, new_encryption_key="wrong-bulk-key") + rows = [ + _prisma_map_row({"server_id": "broken", field: corrupt, "approval_status": "pending_review"}), + _prisma_map_row({"server_id": "healthy", field: encrypted, "approval_status": "active"}), + ] + snapshot = [row.model_dump() for row in rows] + table = SimpleNamespace(find_many=AsyncMock(return_value=rows)) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + result = await reader(prisma, ["broken", "healthy"]) if reader is get_mcp_servers else await reader(prisma) + items = result.items if reader is get_mcp_submissions else result + assert [row.server_id for row in items] == ["healthy"] + assert getattr(items[0], field) == secret + assert [row.model_dump() for row in rows] == snapshot + assert "broken" in caplog.text + assert all(value not in caplog.text for value in ("bulk-sensitive-secret", corrupt, encrypted)) + if reader is get_mcp_submissions: + assert (result.total, result.pending_review, result.active, result.rejected) == (1, 0, 1, 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reader", [get_all_mcp_servers, get_mcp_servers, get_mcp_submissions]) +async def test_bulk_reads_do_not_swallow_unrelated_validation_errors(reader): + from pydantic import ValidationError + + row = _prisma_map_row({"server_id": "invalid", "transport": "unsupported"}) + table = SimpleNamespace(find_many=AsyncMock(return_value=[row])) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + request = reader(prisma, ["invalid"]) if reader is get_mcp_servers else reader(prisma) + with pytest.raises(ValidationError, match="transport"): + await request + + # ── BYOK round-trip ─────────────────────────────────────────────────────────── 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 be4206a1faf..763200c3709 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 @@ -4,6 +4,7 @@ import hashlib import json import time from base64 import urlsafe_b64encode +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -18,6 +19,57 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None): + credential = {"type": "oauth2", "access_token": access_token} + if refresh_token is not None: + credential["refresh_token"] = refresh_token + if expires_in_seconds is not None: + credential["expires_at"] = (datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)).isoformat() + if expires_at is not None: + credential["expires_at"] = expires_at + return credential + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("fields", "egress_has_token"), + [ + (None, False), + ({"access_token": "", "refresh_token": "refresh-token"}, False), + ({}, True), + ({"expires_at": "never"}, True), + ({"expires_in_seconds": 600}, True), + ({"expires_in_seconds": 30}, False), + ({"expires_in_seconds": -300}, False), + ({"expires_in_seconds": -300, "refresh_token": ""}, False), + ({"expires_in_seconds": 30, "refresh_token": "refresh-token"}, True), + ({"expires_in_seconds": -300, "refresh_token": "refresh-token"}, True), + ], +) +async def test_vendor_credential_state_agrees_with_egress_token_resolution(monkeypatch, fields, egress_has_token): + from litellm.proxy._experimental.mcp_server import db as mcp_db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + + monkeypatch.setattr(mcp_db, "MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", 60) + credential = _stored_grant(**fields) if fields is not None else None + read = AsyncMock(return_value=credential) + refresh = AsyncMock(return_value=_stored_grant(access_token="fresh-token", expires_in_seconds=3600)) + prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr(mcp_db, "get_user_oauth_credential", read) + monkeypatch.setattr(mcp_db, "refresh_user_oauth_token", refresh) + + connect = await discoverable_endpoints._vendor_credential_state("user-1", "server-1") + read.assert_awaited_once_with(prisma, "user-1", "server-1") + refresh.assert_not_awaited() + egress = await mcp_db.resolve_valid_user_oauth_token( + user_id="user-1", server=MagicMock(), cred=credential, prisma_client=prisma + ) + + assert (egress is not None) is egress_has_token + assert connect == ("present" if egress_has_token else "absent") + + # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control @pytest.fixture(autouse=True) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 1670370f082..73a52a8d2e8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( aggregate_authorize, aggregate_token, complete_connect_flow, + describe_connect_flow, introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, @@ -230,7 +231,7 @@ async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_co assert location.path == "/ui/connect" params = parse_qs(location.query) handle = params["connect_flow"][0] - assert params["connect_client"] == ["https://claude.ai"] + assert set(params) == {"connect_flow"} set_cookie = response.headers["set-cookie"] assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie assert "HttpOnly" in set_cookie @@ -889,14 +890,60 @@ def _opened_principal(payload): return admitted.principal -async def _finish_connect_page(response): +class _VendorCredential: + def __init__(self, state="present"): + self.calls = [] + self.state = state + + async def __call__(self, user_id, server_id): + self.calls.append((user_id, server_id)) + return self.state + + +class _ServerReachability: + def __init__(self, reachable=True): + self.calls = [] + self.reachable = reachable + + async def __call__(self, user_id, server_id): + self.calls.append((user_id, server_id)) + return self.reachable + + +async def _complete_page(response, scoped_server=None, vendor=None, reachable=None, cache=None, **overrides): + from unittest.mock import patch + handle, cookies = _flow_cookie_from(response) - completed = await complete_connect_flow( - request=_request("/authorize/complete", cookies=cookies, method="POST"), - flow_handle=handle, - session_user_id="u1", - cache=DualCache(), - ) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.return_value = scoped_server + return await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=cache or DualCache(), + lookup_vendor_credential=vendor or _VendorCredential(), + lookup_server_reachability=reachable or _ServerReachability(), + **overrides, + ) + + +async def _describe_page(response, scoped_server=None, vendor=None, reachable=None, session_user_id="u1", cookies=None): + from unittest.mock import patch + + handle, flow_cookies = _flow_cookie_from(response) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.return_value = scoped_server + return await describe_connect_flow( + request=_request("/authorize/flow", cookies=flow_cookies if cookies is None else cookies), + flow_handle=handle, + session_user_id=session_user_id, + lookup_vendor_credential=vendor or _VendorCredential(), + lookup_server_reachability=reachable or _ServerReachability(), + ) + + +async def _finish_connect_page(response, scoped_server=None): + completed = await _complete_page(response, scoped_server=scoped_server) return parse_qs(urlparse(completed.headers["location"]).query)["code"][0] @@ -910,23 +957,43 @@ def _sealed_wire_json(sealed, prefix, debug_key): @pytest.mark.asyncio async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): - """LIT-4917: a per-server RFC 8707 resource naming a gateway-managed oauth2 server - seals that server into the flow. The connect page interlude runs exactly as before - (the scope restricts, it never skips consent), and the code minted at the finish step - and the session pair it redeems for are both scoped.""" + """LIT-4917 plus LIT-7075: a per-server RFC 8707 resource naming a gateway-managed oauth2 + server seals that server into the flow. The connect URL carries only the handle; the page + learns the scoped server and its vendor state from describe_connect_flow, and the finish + step refuses to mint a scoped code until that vendor credential exists, without burning + the flow. The code minted afterwards and the session pair it redeems for are both scoped.""" from unittest.mock import patch client_id = (await _register([REDIRECT_URI]))["client_id"] + github = _scoped_mcp_server() with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + manager.get_mcp_server_by_name.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) assert response.status_code == 303 - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) assert ( _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == "github-id" ) - code = await _finish_connect_page(response) + described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("absent")) + assert json.loads(described.body) == { + "state": "interactive", + "client_origin": "https://claude.ai", + "server_id": "github-id", + "server_name": "github", + "connected": False, + } + cache = DualCache() + premature = await _complete_page(response, scoped_server=github, vendor=_VendorCredential("absent"), cache=cache) + assert premature.status_code == 400 + assert "authorize the requested MCP server" in json.loads(premature.body)["error_description"] + present = _VendorCredential("present") + completed = await _complete_page(response, scoped_server=github, vendor=present, cache=cache) + assert completed.status_code == 303 + assert present.calls == [("u1", "github-id")] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] assert ( _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"] == "github-id" @@ -952,9 +1019,11 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): ) async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, resolves): """Every resource shape outside 'exactly one gateway-managed server' keeps today's flow: - connect page interlude, and NONE of the minted artifacts carry the scope key on the - wire, not the flow cookie, not the code, not the session JWT, so an unscoped flow - started on a new pod completes on a pod whose strict models predate the claim.""" + the generic connect grid (describe names no server, the finish step never consults the + vendor credential), and NONE of the minted + artifacts carry the scope key on the wire, not the flow cookie, not the code, not the + session JWT, so an unscoped flow started on a new pod completes on a pod whose strict + models predate the claim.""" import base64 from unittest.mock import patch @@ -963,10 +1032,18 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, manager.get_mcp_server_by_name.return_value = None if resolves is None else _scoped_mcp_server() response = _scoped_authorize(client_id, resource) assert response.status_code == 303 - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) assert "resource_server_id" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") - code = await _finish_connect_page(response) + vendor = _VendorCredential("absent") + described = await _describe_page(response, vendor=vendor) + assert json.loads(described.body)["state"] == "unscoped" + assert json.loads(described.body)["server_id"] is None + completed = await _complete_page(response, vendor=vendor) + assert vendor.calls == [] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] assert "resource_server_id" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code") token_response = await _redeem(code, client_id) payload = json.loads(token_response.body) @@ -979,19 +1056,150 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, @pytest.mark.asyncio async def test_scoped_authorize_delegate_server_stays_unscoped(): """A delegate-auth oauth2 server is outside the gateway-managed set (its keyless flow is - upstream PKCE via the relay), so a resource naming it never scopes the gateway flow.""" + upstream PKCE via the relay), so a resource naming it never scopes the gateway flow and + never narrows the connect page to it.""" from unittest.mock import patch client_id = (await _register([REDIRECT_URI]))["client_id"] with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True) response = _scoped_authorize(client_id, SCOPED_RESOURCE) - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} + assert json.loads((await _describe_page(response)).body)["server_id"] is None code = await _finish_connect_page(response) token_response = await _redeem(code, client_id) assert _opened_principal(json.loads(token_response.body)).resource_server_id is None +@pytest.mark.asyncio +async def test_m2m_scoped_flow_mints_without_a_user_credential(): + """A client-credentials server is already authorized by its gateway service credential, so + a resource-scoped flow finishes without consulting the per-user vault.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + m2m = _scoped_mcp_server(oauth2_flow="client_credentials") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = m2m + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + vendor = _VendorCredential("unavailable") + described = await _describe_page(response, scoped_server=m2m, vendor=vendor) + assert json.loads(described.body)["state"] == "m2m" + assert json.loads(described.body)["connected"] is True + assert vendor.calls == [] + completed = await _complete_page(response, scoped_server=m2m, vendor=vendor) + assert completed.status_code == 303 + assert vendor.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("oauth2_flow", ["authorization_code", "client_credentials"]) +async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow): + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server(oauth2_flow=oauth2_flow) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + reachable = _ServerReachability(False) + vendor = _VendorCredential("present") + described = await _describe_page(response, scoped_server=server, reachable=reachable, vendor=vendor) + assert json.loads(described.body) == { + "state": "stale", + "client_origin": "https://claude.ai", + "server_id": None, + "server_name": None, + "connected": None, + } + cache = DualCache() + refused = await _complete_page(response, scoped_server=server, reachable=reachable, vendor=vendor, cache=cache) + assert refused.status_code == 400 + assert vendor.calls == [] + assert reachable.calls == [("u1", "github-id"), ("u1", "github-id")] + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + + +@pytest.mark.asyncio +async def test_stale_scoped_flow_remains_distinct_from_unscoped(): + """A server removed after authorize stays a stale scoped flow, so the page cannot offer a + broader unscoped grant or report a misleading Finish action.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + described = await _describe_page(response, scoped_server=None) + assert json.loads(described.body)["state"] == "stale" + assert json.loads(described.body)["connected"] is None + stale = await _complete_page(response, scoped_server=None) + assert stale.status_code == 400 + assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available" + + +@pytest.mark.asyncio +async def test_scoped_flow_deny_and_stale_server_never_need_the_vendor_credential(): + """Cancel is the escape hatch: a scoped user who cannot finish the vendor step still ends + the flow with access_denied and no credential lookup. A scoped server that is no longer + gateway-managed refuses to mint (nothing could serve that code) but also burns nothing.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + github = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = github + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + cache = DualCache() + skipped_reachability = _ServerReachability(False) + stale = await _complete_page(response, scoped_server=None, reachable=skipped_reachability, cache=cache) + assert stale.status_code == 400 + assert skipped_reachability.calls == [] + assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available" + described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("unavailable")) + assert described.status_code == 503 + vendor = _VendorCredential("absent") + deny_reachability = _ServerReachability(False) + denied = await _complete_page( + response, + scoped_server=github, + vendor=vendor, + reachable=deny_reachability, + cache=cache, + decision="deny", + ) + assert denied.status_code == 303 + assert parse_qs(urlparse(denied.headers["location"]).query)["error"] == ["access_denied"] + assert vendor.calls == [] + assert deny_reachability.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "session_user_id, cookies, expected_status, expected_error", + [ + ("u1", {}, 400, "invalid_request"), + ("u1", {"mcp_connect_flow_wrong": "garbage"}, 400, "invalid_request"), + (None, None, 401, "login_required"), + ("u2", None, 403, "access_denied"), + ], +) +async def test_describe_connect_flow_refuses_exactly_like_the_finish_step( + session_user_id, cookies, expected_status, expected_error +): + """The page's read of the flow is gated the same way minting is: the HttpOnly cookie for + that handle must open and the signed-in user must be the sealed one. A lure link with a + made-up handle therefore learns nothing and starts nothing.""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _authorize(client_id, session_user_id="u1") + described = await _describe_page(response, session_user_id=session_user_id, cookies=cookies) + assert described.status_code == expected_status + assert json.loads(described.body)["error"] == expected_error + + @pytest.mark.asyncio async def test_token_rejects_resource_conflicting_with_sealed_scope(): """RFC 8707 section 2.2: redeeming a scoped code (or rotating a scoped refresh token) @@ -1005,7 +1213,7 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope(): with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) - code = await _finish_connect_page(response) + code = await _finish_connect_page(response, scoped_server=github) with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = linear diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py index 76cc235f7eb..36b545ad031 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_env_vars.py @@ -861,6 +861,7 @@ _SALT_KEY = "test-salt-key-for-env-vars-tests-1234" @pytest.fixture def env_vars_salt_key(monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", _SALT_KEY) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"encryption_algorithm": "xsalsa20-poly1305"}) def _mock_env_vars_prisma(row=None): @@ -1518,9 +1519,16 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri assert "s3cr3t-p@ss" not in encrypted_env_vars_str def _prisma_row_with_json_string_env_vars(): - row = MagicMock() - row.env_vars = encrypted_env_vars_str - return row + import json + + from prisma.models import LiteLLM_MCPServerTable + + return LiteLLM_MCPServerTable.model_validate({ + "server_id": "srv-returned", "transport": "http", "mcp_access_groups": [], "allowed_tools": [], + "extra_headers": [], "args": [], "allow_all_keys": False, "available_on_public_internet": True, + "delegate_auth_to_upstream": False, "oauth_passthrough": False, "per_server_oauth_discovery": False, + "is_byok": False, "byok_description": [], "env_vars": json.dumps(encrypted_env_vars_str), + }) mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.create = AsyncMock( @@ -1537,7 +1545,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri touched_by="test-user", ) assert isinstance(created.env_vars, list) - assert created.env_vars[0]["value"] == "s3cr3t-p@ss" + assert created.env_vars[0].value == "s3cr3t-p@ss" + assert created.env_vars[0].name == "DB_PASSWORD" + assert created.env == {} mock_prisma_upd = MagicMock() mock_prisma_upd.db.litellm_mcpservertable.update = AsyncMock( @@ -1549,7 +1559,9 @@ async def test_create_mcp_server_decrypts_env_vars_when_prisma_returns_json_stri touched_by="test-user", ) assert isinstance(updated.env_vars, list) - assert updated.env_vars[0]["value"] == "s3cr3t-p@ss" + assert updated.env_vars[0].value == "s3cr3t-p@ss" + assert updated.env_vars[0].name == "DB_PASSWORD" + assert updated.env == {} def test_reencrypt_global_env_var_values_handles_json_string(env_vars_salt_key): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py index 6bc684af8c3..29e84331f6f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_partial_update.py @@ -12,7 +12,7 @@ import json from unittest.mock import AsyncMock, MagicMock import pytest -from prisma import Json +from prisma import Json, models from litellm.proxy._experimental.mcp_server.db import ( create_mcp_server, @@ -29,8 +29,11 @@ def _credentials_cleared(value) -> bool: def _mock_prisma(): mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable = AsyncMock() - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) - mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=MagicMock()) + row = models.LiteLLM_MCPServerTable.model_construct( + server_id="test-server", transport="http", env={}, env_vars=[] + ) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=row) + mock_prisma.db.litellm_mcpservertable.create = AsyncMock(return_value=row) return mock_prisma diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 086ab854e36..79a2a27bb6d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -2,6 +2,7 @@ import asyncio import contextvars import os from datetime import datetime, timedelta +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1368,6 +1369,295 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): mock_logger.info.assert_any_call("Successfully fetched %s tools total from all MCP servers", 0) +def _denied_scope_manager(known_server_names_to_ids: dict[str, str]) -> MagicMock: + servers = {name: MagicMock(server_id=server_id) for name, server_id in known_server_names_to_ids.items()} + manager = MagicMock() + manager.get_mcp_server_by_name = lambda name, client_ip=None: servers.get(name) + return manager + + +def _scope_resolver(resolved_without_agent: dict[str, str], access_groups: tuple[str, ...] = ()) -> AsyncMock: + async def resolve(user_api_key_auth, mcp_servers, client_ip=None): + if user_api_key_auth is not None and user_api_key_auth.agent_id: + return [] + return [ + SimpleNamespace( + server_id=server_id, + server_name=server_name, + alias=None, + short_prefix=None, + access_groups=list(access_groups), + ) + for server_name, server_id in resolved_without_agent.items() + ] + + return AsyncMock(side_effect=resolve) + + +async def _denied_scoped_list( + user_api_key_auth: UserAPIKeyAuth, + mcp_servers: list[str], + mock_manager: MagicMock, + resolver: AsyncMock, +) -> HTTPException: + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + + with ( + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + resolver, + ), + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + mock_manager, + ), + ): + with pytest.raises(HTTPException) as exc_info: + await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=None, + mcp_servers=mcp_servers, + ) + return exc_info.value + + +@pytest.mark.asyncio +async def test_scoped_list_denied_by_agent_binding_raises_403_naming_agent(): + """The agent-binding veto must raise a 403 naming the agent, never a silent 200 with no tools.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + resolver = _scope_resolver(resolved_without_agent={"github": "srv-github"}) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "MCP server 'github'" in message + assert "agent 'agent-123'" in message + assert "mcp_servers" in message + rerun_auth = resolver.await_args_list[1].kwargs["user_api_key_auth"] + assert rerun_auth.agent_id is None + assert rerun_auth.user_id == "test_user" + assert resolver.await_args_list[1].kwargs["mcp_servers"] == ["github"] + + +@pytest.mark.asyncio +async def test_empty_scope_lists_nothing_instead_of_raising_a_nameless_denial(): + """An empty ``x-mcp-servers`` header scopes to no servers; that is an empty listing, not a 403.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + resolver = AsyncMock(return_value=[]) + + with ( + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + resolver, + ), + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager", + _denied_scope_manager({"github": "srv-github"}), + ), + ): + listing = await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, mcp_auth_header=None, mcp_servers=[] + ) + + assert listing.tools == [] + resolver.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scoped_list_denied_for_non_agent_key_raises_generic_403(): + """A denial for a key with no agent binding stays generic and skips the agent-stripped rerun.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + resolver = AsyncMock(return_value=[]) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "github" in message + assert "agent" not in message + resolver.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scoped_list_unknown_name_raises_same_generic_403_as_unauthorized(): + """Unknown and registered-but-unauthorized names raise byte-identical generic 403s, so a + caller cannot probe which server names exist.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + unknown = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({}), _scope_resolver(resolved_without_agent={}) + ) + unauthorized = await _denied_scoped_list( + user_api_key_auth, + ["github"], + _denied_scope_manager({"github": "srv-github"}), + _scope_resolver(resolved_without_agent={}), + ) + + assert unknown.status_code == unauthorized.status_code == 403 + assert unknown.detail["error"] == unauthorized.detail["error"] + assert "github" in unknown.detail["error"] + assert "agent" not in unknown.detail["error"] + + +@pytest.mark.asyncio +async def test_scoped_list_access_group_vetoed_by_agent_names_agent_and_group(): + """An access-group scope vetoed by the agent binding raises the 403 naming the agent and the + group instead of the silent empty list.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["prod-group"], + _denied_scope_manager({}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "access group 'prod-group'" in message + assert "agent 'agent-123'" in message + assert "mcp_access_groups" in message + + +@pytest.mark.asyncio +async def test_scoped_list_mixed_unknown_and_vetoed_group_names_the_group_that_resolved(): + """With an unknown name ahead of the agent-vetoed group in the scope, the 403 must name the group + whose servers the key can reach, never the unknown name, or the admin is told to grant a group + that does not exist.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["no-such-group", "prod-group"], + _denied_scope_manager({}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}, access_groups=("prod-group",)), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "access group 'prod-group'" in message + assert "no-such-group" not in message + assert "agent 'agent-123'" in message + + +@pytest.mark.asyncio +async def test_scoped_list_agent_key_denied_by_key_grants_raises_generic_403(): + """When the agent-stripped rerun still resolves nothing, the 403 stays generic instead of + blaming the agent binding.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + resolver = _scope_resolver(resolved_without_agent={}) + + denial = await _denied_scoped_list( + user_api_key_auth, ["github"], _denied_scope_manager({"github": "srv-github"}), resolver + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "github" in message + assert "agent" not in message + assert resolver.await_count == 2 + + +@pytest.mark.asyncio +async def test_scoped_list_agent_veto_attributed_for_differently_cased_server_name(): + """The scope filter matches `/mcp/GitHub` to a server named `github` case-insensitively, so the + agent-attributed 403 must match the same way instead of falling back to the generic denial.""" + pytest.importorskip("litellm.proxy._experimental.mcp_server.server") + + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user", agent_id="agent-123") + + denial = await _denied_scoped_list( + user_api_key_auth, + ["GitHub"], + _denied_scope_manager({"github": "srv-github"}), + _scope_resolver(resolved_without_agent={"github": "srv-github"}), + ) + + assert denial.status_code == 403 + message = denial.detail["error"] + assert "MCP server 'GitHub'" in message + assert "agent 'agent-123'" in message + + +@pytest.mark.asyncio +async def test_handle_list_tools_converts_permission_httpexception_to_mcp_error(): + """The MCP protocol handler surfaces a permission HTTPException as a clean JSON-RPC error + (McpError, INVALID_REQUEST) carrying the denial message, instead of a raw 500.""" + try: + from litellm.proxy._experimental.mcp_server.server import handle_list_tools + except ImportError: + pytest.skip("MCP server not available") + + from mcp.shared.exceptions import McpError + from mcp.types import INVALID_REQUEST + + denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'" + denial = HTTPException(status_code=403, detail={"error": denial_message}) + + with ( + patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + ), + patch( # test-quality-ok: the listing helper is the handler's only collaborator; the suite's seam + "litellm.proxy._experimental.mcp_server.server._list_mcp_tools", + new=AsyncMock(side_effect=denial), + ), + ): + with pytest.raises(McpError) as exc_info: + await handle_list_tools() + + assert exc_info.value.error.code == INVALID_REQUEST + assert exc_info.value.error.message == denial_message + + +@pytest.mark.asyncio +async def test_mcp_server_tool_call_renders_denial_message_not_detail_dict(): + try: + from litellm.proxy._experimental.mcp_server.server import mcp_server_tool_call + except ImportError: + pytest.skip("MCP server not available") + + denial_message = "MCP server 'github' is not available to this key: the key is bound to agent 'agent-123'" + denial = HTTPException(status_code=403, detail={"error": denial_message}) + + with ( + patch( # test-quality-ok: the protocol handler reads auth from module context; no injection seam + "litellm.proxy._experimental.mcp_server.server.get_or_extract_auth_context", + new=AsyncMock(return_value=(None, None, None, None, None, None, None)), + ), + patch( # test-quality-ok: the tool-call helper is the handler's only collaborator; the suite's seam + "litellm.proxy._experimental.mcp_server.server.call_mcp_tool", + new=AsyncMock(side_effect=denial), + ), + ): + result = await mcp_server_tool_call("github-search_issues", {}) + + assert result.isError is True + assert result.content[0].text == f"Error: {denial_message}" + + @pytest.mark.asyncio async def test_mcp_server_tool_call_body_with_none_arguments(): """Test that proxy_server_request body handles None arguments correctly""" @@ -3518,6 +3808,35 @@ async def test_call_mcp_tool_user_unauthorized_access(): assert "User not allowed to call this tool" in exc_info.value.detail +@pytest.mark.asyncio +async def test_call_mcp_tool_scoped_denial_names_the_binding_agent(): + from litellm.proxy._experimental.mcp_server.server import call_mcp_tool + + agent_bound_key = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-123") + + with ( + patch( # test-quality-ok: the server registry is a module-level singleton; the suite's only seam + "litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_allowed_mcp_servers", + AsyncMock(return_value=[]), + ), + patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + _scope_resolver({"github": "srv-github"}), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await call_mcp_tool( + name="github-search_issues", + arguments={}, + user_api_key_auth=agent_bound_key, + mcp_servers=["github"], + ) + + assert exc_info.value.status_code == 403 + assert "MCP server 'github'" in exc_info.value.detail["error"] + assert "agent 'agent-123'" in exc_info.value.detail["error"] + + @pytest.mark.asyncio async def test_call_mcp_tool_unauthorized_403_does_not_leak_server_credentials(): """Regression for LIT-4703 / GH #29936. 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 d19363d3b5f..46fef83092d 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 @@ -905,6 +905,68 @@ class TestMCPServerManager: assert retry_slot is not None assert retry_slot.generation > old_generation + @pytest.mark.asyncio + @pytest.mark.parametrize("corrupt_column", ("static_headers", "env")) + async def test_database_reload_drops_cached_server_whose_secret_map_stops_decoding( + self, monkeypatch, caplog, corrupt_column + ): + from types import SimpleNamespace + + from litellm.proxy import proxy_server + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_secret_map + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-reload-secret-map-salt") + monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": "aes-256-gcm"}) + headers = {"Authorization": "Bearer dummy-header-secret-4f1c"} + env = {"UPSTREAM_TOKEN": "dummy-env-secret-9a2b"} + stamp = datetime.now() + cached = MCPServer( + server_id="cached-server", + name="cached_server", + url="https://up.example.com/mcp", + transport=MCPTransport.http, + static_headers=dict(headers), + env=dict(env), + updated_at=stamp, + ) + manager = MCPServerManager() + manager.registry[cached.server_id] = cached + stored = {"static_headers": encrypt_secret_map(headers), "env": encrypt_secret_map(env)} + corrupted = {**stored, corrupt_column: stored[corrupt_column][:-6] + 'AAAAA"'} + + def _row(server_id, maps): + row = MagicMock() + row.server_id = server_id + row.alias = server_id + row.model_dump.return_value = { + "server_id": server_id, + "alias": server_id, + "server_name": server_id, + "url": "https://up.example.com/mcp", + "transport": MCPTransport.http, + "updated_at": stamp, + **maps, + } + return row + + table = SimpleNamespace( + find_many=AsyncMock(return_value=[_row(cached.server_id, corrupted), _row("healthy-sibling", stored)]) + ) + prisma = SimpleNamespace(db=SimpleNamespace(litellm_mcpservertable=table)) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + await manager.reload_servers_from_database() + + assert set(manager.registry) == {"healthy-sibling"} + sibling = manager.registry["healthy-sibling"] + assert dict(sibling.static_headers) == headers + assert dict(sibling.env) == env + logged = "\n".join(caplog.messages) + assert cached.server_id in logged + for secret in (*headers.values(), *env.values(), *stored.values(), corrupted[corrupt_column]): + assert secret not in logged + @pytest.mark.asyncio async def test_lazy_oauth_discovery_preserves_manual_authorization_url_gate(self): with patch.dict(os.environ, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "false"}): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py index f6b61c1d9f7..e814425c9a2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_sigv4_auth.py @@ -15,6 +15,11 @@ import httpx from litellm.experimental_mcp_client.client import MCPSigV4Auth, MCPClient from litellm.types.mcp import MCPAuth, MCPTransport +from prisma import models + + +def _updated_row() -> models.LiteLLM_MCPServerTable: + return models.LiteLLM_MCPServerTable.model_construct(server_id="test-server", transport="http", env={}, env_vars=[]) class TestMCPSigV4Auth: @@ -600,7 +605,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -639,7 +644,7 @@ class TestCredentialMergeOnUpdate: from litellm.proxy._types import UpdateMCPServerRequest mock_prisma = MagicMock() - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -667,7 +672,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -709,7 +714,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -752,7 +757,7 @@ class TestCredentialMergeOnUpdate: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", @@ -1083,7 +1088,7 @@ class TestAuthTypeSwitchClearsCredentials: mock_prisma = MagicMock() mock_prisma.db.litellm_mcpservertable.find_unique = AsyncMock(return_value=existing_record) - mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=MagicMock()) + mock_prisma.db.litellm_mcpservertable.update = AsyncMock(return_value=_updated_row()) data = UpdateMCPServerRequest( server_id="test-server", 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..b8935d07774 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 @@ -11,6 +11,7 @@ Covers: import json from collections.abc import Sequence +from types import SimpleNamespace from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -24,6 +25,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 +274,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 +332,7 @@ class TestGetVirtualToolDefinitions: MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, } @@ -364,7 +367,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 +745,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 +815,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 +833,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 +850,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 +1238,7 @@ class TestHandleListToolsVirtual: MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME, AGENT_SEARCH_TOOL_NAME, + SKILL_SEARCH_TOOL_NAME, } @@ -1229,3 +1272,37 @@ class TestMcpServerToolCallErrorHandling: assert result.isError is True assert "User not allowed to call this tool" in result.content[0].text + + +@pytest.mark.asyncio +async def test_handle_mcp_tool_call_scoped_denial_names_the_binding_agent() -> None: + from fastapi import HTTPException + + from litellm.proxy._experimental.mcp_server.tool_search import handle_mcp_tool_call + + agent_bound_key = UserAPIKeyAuth(api_key="test_key", agent_id="agent-123") + + async def resolve(user_api_key_auth, mcp_servers, client_ip=None): + if user_api_key_auth.agent_id: + return [] + return [ + SimpleNamespace( + server_id="srv-github", server_name="github", alias=None, short_prefix=None, access_groups=[] + ) + ] + + with patch( # test-quality-ok: the permission resolver is a module-level function; the suite's only seam + "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers", + new=AsyncMock(side_effect=resolve), + ): + with pytest.raises(HTTPException) as exc_info: + await handle_mcp_tool_call( + tool_name="github-create_issue", + arguments={}, + user_api_key_dict=agent_bound_key, + mcp_servers=["github"], + ) + + assert exc_info.value.status_code == 403 + assert "MCP server 'github'" in exc_info.value.detail["error"] + assert "agent 'agent-123'" in exc_info.value.detail["error"] 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/_experimental/mcp_server/test_semantic_tool_filter.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py index bf0df17fafb..f0b4e94f72f 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_semantic_tool_filter.py @@ -2216,3 +2216,26 @@ async def test_top_k_above_router_default_is_respected(): assert len(filtered) == 6 print("✅ Configured top_k above the semantic-router default of 5 is honored") + + +@pytest.mark.asyncio +async def test_semantic_filter_hook_narrows_only_references_the_gateway_serves(): + from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + + gateway_reference = {"type": "mcp", "server_url": "litellm_proxy", "require_approval": "never"} + external_tool = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "allowed_tools": ["zapier_send_email"], + } + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + narrowed = await SemanticToolFilterHook._narrow_mcp_references( + [gateway_reference, external_tool], ["srv-tool_1"], served_names=served_names + ) + + assert narrowed == [{**gateway_reference, "allowed_tools": ["srv-tool_1"]}, external_tool] 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..aaf630ad29b 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -3,6 +3,7 @@ Unit tests for auth_utils functions related to rate limiting and customer ID ext """ import base64 +import logging from typing import Optional from unittest.mock import MagicMock, patch @@ -15,6 +16,7 @@ from litellm.proxy.auth.auth_utils import ( abbreviate_api_key, check_complete_credentials, custom_auth_common_checks_warning, + log_once_if_budget_reservation_disabled, warn_once_if_custom_auth_skips_common_checks, get_end_user_id_from_request_body, get_key_mcp_rpm_limit, @@ -101,6 +103,41 @@ class TestWarnOnceIfCustomAuthSkipsCommonChecks: assert logger.warning.call_count == 0 +class TestLogOnceIfBudgetReservationDisabled: + @pytest.fixture(autouse=True) + def _reset_sentinel(self, monkeypatch): + monkeypatch.setattr( + "litellm.constants.budget_reservation_disabled_info_emitted", + False, + ) + + def test_logs_info_only_once_when_enabled(self, caplog): + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + log_once_if_budget_reservation_disabled(disabled=False) + assert not any( + "disable_budget_reservation is enabled" in record.message + for record in caplog.records + ) + for _ in range(3): + log_once_if_budget_reservation_disabled(disabled=True) + + records = [ + record + for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert len(records) == 1 + assert records[0].levelno == logging.INFO + + def test_logs_to_injected_logger_only_once(self): + logger = MagicMock() + log_once_if_budget_reservation_disabled(disabled=False, logger=logger) + for _ in range(3): + log_once_if_budget_reservation_disabled(disabled=True, logger=logger) + assert logger.info.call_count == 1 + assert "disable_budget_reservation is enabled" in logger.info.call_args[0][0] + + class TestGetKeyModelRpmLimit: """Tests for get_key_model_rpm_limit function.""" @@ -428,6 +465,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_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/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index d44f96d95bf..541aeabcbcd 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -1,5 +1,6 @@ import asyncio import json +import logging from contextlib import contextmanager from datetime import datetime, timedelta from types import SimpleNamespace @@ -146,6 +147,35 @@ async def test_disable_budget_reservation_skips_reservation(): assert user_api_key_auth_obj.budget_reservation is None +@pytest.mark.asyncio +async def test_disable_budget_reservation_does_not_log_per_request(caplog): + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + for _ in range(3): + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings={"disable_budget_reservation": True}, + ) + + records = [ + record + for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert records == [] + assert user_api_key_auth_obj.budget_reservation is None + + @pytest.mark.asyncio async def test_budget_reservation_runs_when_not_disabled(): """Control for #27639: with the flag absent, the reservation still runs and is stored.""" diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index 3bb3f75b9b8..2a0ebf9492a 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -824,7 +824,7 @@ class _StopFailingSubscriber(ConfigSyncSubscriber): raise RuntimeError("stop failed") -async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> None: +async def test_proxy_config_subscriber_resyncs_deployments_only() -> None: from litellm.proxy.proxy_server import ProxyConfig cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) @@ -852,10 +852,7 @@ async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> await callback() await config.stop_config_sync_subscriber() - assert calls == [ - ("add_deployment", prisma_client, proxy_logging_obj), - ("get_credentials", prisma_client, None), - ] + assert calls == [("add_deployment", prisma_client, proxy_logging_obj)] assert config.config_sync_subscriber is None assert subscriber._task is None diff --git a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py index f0fbdea4e85..6f7c20166c5 100644 --- a/tests/test_litellm/proxy/common_utils/test_registry_read_through.py +++ b/tests/test_litellm/proxy/common_utils/test_registry_read_through.py @@ -393,6 +393,51 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc assert not proxy_server.MODEL_RECONCILE_LOCK.locked() +@pytest.mark.asyncio +async def test_resync_model_deployments_loads_db_credentials_before_reconciling_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unittest.mock import AsyncMock, MagicMock + + import litellm + import litellm.proxy.proxy_server as proxy_server + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments + from litellm.types.utils import CredentialItem + + rows: Final = [MagicMock()] + prisma_client: Final = MagicMock() + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows) + router: Final = MagicMock() + router.get_model_list.return_value = [] + installed: Final = MagicMock() + + async def load_credentials_from_db(prisma_client: object) -> None: + CredentialAccessor.upsert_credentials( + [ + CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db"}, + credential_info={}, + ) + ] + ) + + def install_models(db_models: object) -> None: + installed(db_models=db_models, credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(litellm, "credential_list", []) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "store_model_in_db", True) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "llm_model_list", None) + monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", load_credentials_from_db) + monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install_models) + + assert await _resync_model_deployments("model-created-on-a-sibling-replica") is True + installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"}) + + @pytest.mark.asyncio async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch): from unittest.mock import AsyncMock, MagicMock diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py index ca4d62737b6..54e0aa74a25 100644 --- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py +++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py @@ -23,7 +23,7 @@ from litellm.proxy.common_utils.scheduled_job_stagger import ( ) OPERATOR_CRON_JOB_ID = "spend_log_cleanup_job" -SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "get_credentials_job", "add_deployment_job") +SHARED_INTERVAL_JOB_IDS = ("periodic_reload_job", "add_deployment_job") async def _noop() -> None: ... diff --git a/tests/test_litellm/proxy/db/conftest.py b/tests/test_litellm/proxy/db/conftest.py index 6f67b91ac1d..03d7ea81257 100644 --- a/tests/test_litellm/proxy/db/conftest.py +++ b/tests/test_litellm/proxy/db/conftest.py @@ -1,5 +1,11 @@ +import json import os +import signal +import sys +import time from collections.abc import Generator +from dataclasses import dataclass +from pathlib import Path from typing import Optional import pytest @@ -75,3 +81,76 @@ def reset_entra_token_provider_cache() -> Generator[None, None, None]: def unset_database_url(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DATABASE_URL", "about-to-be-unset") monkeypatch.delenv("DATABASE_URL") + + +FAKE_PRISMA_CLI = """#!{python} +import json +import os +import pathlib +import subprocess +import sys +import time + +calls_file = pathlib.Path(os.environ["FAKE_PRISMA_CALLS"]) +earlier_calls = calls_file.read_text().splitlines() if calls_file.exists() else [] +with calls_file.open("a") as log: + print(json.dumps(sys.argv[1:]), file=log) +if not earlier_calls and os.environ.get("FAKE_PRISMA_HANG_FIRST"): + grandchild = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(600)"]) + pathlib.Path(os.environ["FAKE_PRISMA_GRANDCHILD_PIDFILE"]).write_text(str(grandchild.pid)) + time.sleep(600) +sys.exit(0) +""" + + +@dataclass(frozen=True, slots=True) +class FakePrismaCli: + """A stand-in `prisma` on PATH, recording every invocation. + + With FAKE_PRISMA_HANG_FIRST set it hangs on its first call from a process tree + of its own, the way the real CLI wraps Node around a Rust schema engine, so a + timeout that kills only the direct child leaves the rest of that tree running. + """ + + calls_file: Path + grandchild_pidfile: Path + + @property + def calls(self) -> list[list[str]]: + if not self.calls_file.exists(): + return [] + return [json.loads(line) for line in self.calls_file.read_text().splitlines()] + + def grandchild_is_gone(self, within_seconds: float) -> bool: + deadline = time.monotonic() + within_seconds + while time.monotonic() < deadline: + try: + os.kill(int(self.grandchild_pidfile.read_text()), 0) + except ProcessLookupError: + return True + time.sleep(0.05) + return False + + +@pytest.fixture +def fake_prisma_cli(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[FakePrismaCli, None, None]: + bin_dir = tmp_path / "fakebin" + bin_dir.mkdir() + script = bin_dir / "prisma" + script.write_text(FAKE_PRISMA_CLI.format(python=sys.executable)) + script.chmod(0o755) + cli = FakePrismaCli( + calls_file=tmp_path / "calls.jsonl", + grandchild_pidfile=tmp_path / "grandchild.pid", + ) + monkeypatch.setenv("PATH", f"{bin_dir}{os.pathsep}{os.environ['PATH']}") + monkeypatch.setenv("FAKE_PRISMA_CALLS", str(cli.calls_file)) + monkeypatch.setenv("FAKE_PRISMA_GRANDCHILD_PIDFILE", str(cli.grandchild_pidfile)) + monkeypatch.setenv("LITELLM_PRISMA_COMMAND_TIMEOUT", "1") + monkeypatch.delenv("FAKE_PRISMA_HANG_FIRST", raising=False) + yield cli + if cli.grandchild_pidfile.exists(): + try: + os.kill(int(cli.grandchild_pidfile.read_text()), signal.SIGKILL) + except ProcessLookupError: + pass diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index 2ed4f843711..4507892bd0f 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -9,13 +9,15 @@ request-time transaction builder and the flush contract with an injected fake cl import asyncio import json from datetime import datetime +from types import SimpleNamespace +from typing import Final import httpx import pytest from litellm.proxy.db.autorouter_session_rollup import ( - AutoRouterTurnTransaction, UPSERT_AUTOROUTER_SESSION_SQL, + AutoRouterTurnTransaction, build_autorouter_turn_transaction, flush_autorouter_turn_transactions, ) @@ -70,6 +72,7 @@ class TestBuildTransaction: total_tokens=100, spend=0.01, saved_spend=0.02, + classifier_cost=0.0, covered=True, cache_hit=True, cache_ttl_seconds=300, @@ -111,6 +114,9 @@ class TestBuildTransaction: folded once into the turn that paid for it (GH #38816).""" transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.005})) assert transaction is not None and transaction.spend == pytest.approx(0.015) + assert transaction.classifier_cost == 0.005 + assert transaction.spend - transaction.classifier_cost == pytest.approx(0.01) + assert transaction.saved_spend == 0.02 @pytest.mark.parametrize( "decision_extra", [{}, {"classifier_cost": 0.0}, {"classifier_cost": "bogus"}, {"classifier_cost": True}] @@ -118,6 +124,8 @@ class TestBuildTransaction: def test_an_unpriced_classifier_leaves_the_spend_alone(self, decision_extra: dict): transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, **decision_extra})) assert transaction is not None and transaction.spend == pytest.approx(0.01) + assert transaction.classifier_cost == 0.0 + assert transaction.saved_spend == 0.02 def test_every_turn_carries_its_own_classifier_charge(self): first = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "classifier_cost": 0.005})) @@ -127,6 +135,7 @@ class TestBuildTransaction: ) assert first is not None and first.spend == pytest.approx(0.015) assert second is not None and second.spend == pytest.approx(0.027) + assert (first.classifier_cost, second.classifier_cost) == (0.005, 0.007) def test_router_name_falls_back_to_the_payload_model_group(self): transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"})) @@ -217,6 +226,7 @@ def _transaction( total_tokens=100, spend=0.01, saved_spend=0.02, + classifier_cost=0.005, covered=True, cache_hit=False, cache_ttl_seconds=None, @@ -249,6 +259,7 @@ class TestFlush: 100, 0.01, 0.02, + 0.005, 1, 0, None, @@ -279,26 +290,31 @@ class TestFlush: class TestEnqueueSeam: @pytest.mark.asyncio - async def test_update_database_seam_enqueues_only_auto_routed_success(self, monkeypatch: pytest.MonkeyPatch): + @pytest.mark.parametrize("classifier_cost", [0.005, 0.0, None]) + async def test_update_database_seam_enqueues_only_auto_routed_success(self, classifier_cost: float | None): from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter - from litellm.proxy.utils import PrismaClient - monkeypatch.setattr(PrismaClient, "autorouter_turn_transactions", []) - writer = DBSpendUpdateWriter() - fake_prisma = type("P", (), {})() - fake_prisma._autorouter_turn_transactions_lock = asyncio.Lock() - fake_prisma.autorouter_turn_transactions = [] + writer: Final = DBSpendUpdateWriter() + fake_prisma: Final = SimpleNamespace( + _autorouter_turn_transactions_lock=asyncio.Lock(), autorouter_turn_transactions=[] + ) + metadata: Final = _metadata( + routing_decision={**ROUTING_DECISION, "classifier_cost": classifier_cost}, autorouter_savings=-0.003 + ) + for payload in ( + _payload(metadata=json.dumps(metadata)), + _payload(metadata=json.dumps({"usage_object": {"prompt_tokens": 9}})), + _payload(status="failure", metadata=json.dumps(metadata)), + _payload(metadata=json.dumps({**metadata, "internal_call_origin": "autorouter_classifier"})), + ): + await writer._enqueue_autorouter_turn_transaction(payload=payload, prisma_client=fake_prisma) - routed = _payload() - routed["metadata"] = json.dumps(_metadata()) - await writer._enqueue_autorouter_turn_transaction(payload=routed, prisma_client=fake_prisma) - - plain = _payload() - plain["metadata"] = json.dumps({"usage_object": {"prompt_tokens": 9}}) - await writer._enqueue_autorouter_turn_transaction(payload=plain, prisma_client=fake_prisma) - - assert [t.router_name for t in fake_prisma.autorouter_turn_transactions] == ["live-auto"] - assert fake_prisma.autorouter_turn_transactions[0].saved_spend == 0.0 + assert len(fake_prisma.autorouter_turn_transactions) == 1 + transaction: Final = fake_prisma.autorouter_turn_transactions[0] + assert transaction.router_name == "live-auto" + assert transaction.spend == pytest.approx(0.01 + (classifier_cost or 0.0)) + assert transaction.classifier_cost == (classifier_cost or 0.0) + assert transaction.saved_spend == -0.003 def test_every_drain_trigger_reads_the_one_queue_census_owner(): diff --git a/tests/test_litellm/proxy/db/test_check_migration.py b/tests/test_litellm/proxy/db/test_check_migration.py index 9e2f6a1089c..74a6841cb23 100644 --- a/tests/test_litellm/proxy/db/test_check_migration.py +++ b/tests/test_litellm/proxy/db/test_check_migration.py @@ -37,3 +37,35 @@ def test_check_migration_out_of_sync(mocker): check_migration.verbose_logger.exception.assert_called_once() actual_message = check_migration.verbose_logger.exception.call_args[0][0] assert "prisma schema out of sync with db" in actual_message + + +@pytest.mark.timeout(30) +def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fake_prisma_cli, monkeypatch): + """ + `prisma migrate diff` ran unbounded, so a database that never answers hung boot + before uvicorn ever started, and interrupting the proxy orphaned the schema engine. + """ + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) + assert fake_prisma_cli.calls == [ + ["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x", + "--to-schema-datamodel", "./schema.prisma", "--script"] + ] + assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch): + """ + Boot calls this helper directly, so an ImportError here takes the proxy down before + uvicorn starts. An install without the runner must lose the diagnostic, not the proxy. + """ + import sys + + from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper + + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, []) diff --git a/tests/test_litellm/proxy/db/test_prisma_client.py b/tests/test_litellm/proxy/db/test_prisma_client.py index f0983d6bf62..7f43e483557 100644 --- a/tests/test_litellm/proxy/db/test_prisma_client.py +++ b/tests/test_litellm/proxy/db/test_prisma_client.py @@ -10,7 +10,7 @@ from fastapi.testclient import TestClient -from litellm.proxy.db.prisma_client import PrismaWrapper, should_update_prisma_schema +from litellm.proxy.db.prisma_client import PrismaManager, PrismaWrapper, should_update_prisma_schema @pytest.fixture(autouse=True) @@ -193,7 +193,10 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client( mock_new_prisma.connect.assert_awaited_once() -def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): +DB_PUSH_ARGV = ["db", "push", "--accept-data-loss", "--skip-generate"] + + +def test_db_push_applies_replica_identity_full_when_requested(monkeypatch, fake_prisma_cli, unset_database_url): """`prisma db push` bypasses litellm-proxy-extras, so it needs its own call into the opt-in REPLICA IDENTITY FULL step.""" from litellm.proxy.db.prisma_client import PrismaManager @@ -208,14 +211,13 @@ def test_db_push_applies_replica_identity_full_when_requested(monkeypatch): staticmethod(lambda: applied.append(True)), ) - with patch("litellm.proxy.db.prisma_client.subprocess.run") as mock_run: - assert PrismaManager.setup_database(use_migrate=False) is True + assert PrismaManager.setup_database(use_migrate=False) is True - assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV] assert applied == [True] -def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch): +def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch, fake_prisma_cli, unset_database_url): """A doc-partitioned LiteLLM_SpendLogs makes `prisma db push` rewrite the primary key back to ("request_id"), which Postgres rejects; the guard must fail fast with guidance instead of running the push.""" @@ -228,29 +230,23 @@ def test_db_push_is_rejected_when_spend_logs_is_partitioned(monkeypatch): monkeypatch.setattr( ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: True) ) - with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, asserted never reached - "litellm.proxy.db.prisma_client.subprocess.run" - ) as mock_run: - with pytest.raises(RuntimeError) as err: - PrismaManager.setup_database(use_migrate=False) + with pytest.raises(RuntimeError) as err: + PrismaManager.setup_database(use_migrate=False) assert str(err.value) == PARTITIONED_SPEND_LOGS_PUSH_ERROR - mock_run.assert_not_called() + assert fake_prisma_cli.calls == [] -def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch): +def test_db_push_proceeds_when_spend_logs_is_not_partitioned(monkeypatch, fake_prisma_cli, unset_database_url): from litellm.proxy.db.prisma_client import PrismaManager from litellm_proxy_extras.utils import ProxyExtrasDBManager monkeypatch.setattr( ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False) ) - with patch( # test-quality-ok: subprocess.run is the external prisma CLI boundary, not SDK logic - "litellm.proxy.db.prisma_client.subprocess.run" - ) as mock_run: - assert PrismaManager.setup_database(use_migrate=False) is True + assert PrismaManager.setup_database(use_migrate=False) is True - assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"] + assert fake_prisma_cli.calls == [DB_PUSH_ARGV] def _entra_jwt(expires_in_seconds: int) -> str: @@ -377,3 +373,30 @@ def test_minting_without_the_database_env_vars_names_them(azure_env, monkeypatch with pytest.raises(RuntimeError, match="DATABASE_HOST"): wrapper.get_rds_iam_token() + + +@pytest.mark.timeout(45) +def test_db_push_timeout_takes_its_process_tree_with_it(fake_prisma_cli, unset_database_url, monkeypatch): + """ + A timed-out `db push` used to leave Node and the schema engine writing the schema, + so the next attempt pushed into a database the abandoned one was still mutating. + """ + monkeypatch.delenv("LITELLM_SET_REPLICA_IDENTITY_FULL", raising=False) + monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1") + + assert PrismaManager.setup_database(use_migrate=False) is True + assert fake_prisma_cli.calls == [DB_PUSH_ARGV, DB_PUSH_ARGV] + assert fake_prisma_cli.grandchild_is_gone(within_seconds=5) + + +def test_db_push_without_the_prisma_runner_fails_the_migration_instead_of_crashing_boot( + fake_prisma_cli, unset_database_url, monkeypatch +): + """ + An ImportError out of setup_database escapes the caller's RuntimeError handler and + kills boot, bypassing the operator's enforce_prisma_migration_check choice. + """ + monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None) + + assert PrismaManager.setup_database(use_migrate=False) is False + assert fake_prisma_cli.calls == [] 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/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 0dbd4591ac9..87a1b84acc5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -3,6 +3,7 @@ Unit tests for Tool Permission Guardrail (OpenAI tool_calls semantics) """ import json +import logging import re from unittest.mock import patch @@ -520,6 +521,80 @@ class TestToolPermissionGuardrail: ) assert excinfo.value.status_code == 400 + @pytest.mark.asyncio + async def test_async_pre_call_hook_without_tools_logs_skip_at_debug(self, caplog): + data = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + result = await self.guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + assert result is data + skip_levels = [r.levelno for r in caplog.records if "No tools or functions in data" in r.getMessage()] + assert skip_levels == [logging.DEBUG] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + @pytest.mark.asyncio + async def test_async_pre_call_hook_denied_tool_logs_at_info(self, caplog): + data = {"tools": [{"type": "function", "function": {"name": "Read"}}]} + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException): + await self.guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + denied_levels = [ + r.levelno + for r in caplog.records + if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'" + ] + assert denied_levels == [logging.INFO] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + @pytest.mark.asyncio + async def test_async_post_call_success_hook_denied_tool_logs_at_info(self, caplog): + tool_call = {"function": {"name": "Read", "arguments": "{}"}, "type": "function"} + response = ModelResponse(choices=[Choices(message={"tool_calls": [tool_call]})]) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self.guardrail.async_post_call_success_hook( + data={"guardrails": ["test-tool-permission"]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + denied_levels = [ + r.levelno + for r in caplog.records + if r.getMessage() == "Tool Permission Guardrail: Tool 'Read' denied by rule 'deny_read'" + ] + assert denied_levels == [logging.INFO] + assert [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] == [] + + def test_parse_tool_call_arguments_malformed_json_logs_warning(self, caplog): + tool_call = ChatCompletionMessageToolCall(function={"name": "Bash", "arguments": "{not json"}, id="call_1") + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + parsed, error = self.guardrail._parse_tool_call_arguments(tool_call) + + assert parsed is None + assert error == "arguments could not be parsed" + warning_messages = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert len(warning_messages) == 1 + assert warning_messages[0].startswith("Tool Permission Guardrail: Failed to decode arguments for tool Bash") + @pytest.mark.asyncio async def test_async_pre_call_hook_blocks_legacy_functions(self): data = { 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/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 21f8d985f22..dc18e0f7d4a 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 @@ -10,7 +10,6 @@ import pytest from fastapi import HTTPException from pydantic import ValidationError - from litellm.proxy._types import ( LitellmUserRoles, ProxyErrorTypes, @@ -21,11 +20,11 @@ from litellm.proxy.management_endpoints.auto_router_endpoints import ( preview_auto_router_routing, ) from litellm.router import Router -from litellm.types.utils import Choices, Message, ModelResponse from litellm.types.management_endpoints.auto_router_endpoints import ( AutoRouterBenchmarksResponse, AutoRouterRoutingTestRequest, ) +from litellm.types.utils import Choices, Message, ModelResponse ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test", user_id="admin") @@ -489,6 +488,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 +503,7 @@ class TestAutoRouterBenchmarks: user_api_key_dict=ADMIN, start_date="2026-07-01", end_date="2026-08-01", + api_key=api_key, ) ROW = _SessionAggRow( @@ -527,6 +528,8 @@ class TestAutoRouterBenchmarks: total_tokens=4000, spend=10.0, saved_spend=30.0, + classifier_cost=0.4, + classifier_cost_recorded_turns=40, session_seconds=400.0, ) @@ -565,6 +568,7 @@ class TestAutoRouterBenchmarks: totals = _benchmark_totals(losing) assert totals.baseline_spend == 5.0 assert totals.saved_pct == -100.0 + assert totals.classifier_cost == 0.4 def test_an_empty_window_folds_to_zeros(self): from litellm.proxy.management_endpoints.auto_router_endpoints import ( @@ -577,6 +581,7 @@ class TestAutoRouterBenchmarks: assert totals.turns == 0 assert totals.saved_pct == 0.0 assert totals.cache.hit_rate_pct == 0.0 + assert totals.classifier_cost == 0.0 def test_totals_sum_counters_across_groups_before_deriving_ratios(self): from litellm.proxy.management_endpoints.auto_router_endpoints import ( @@ -652,11 +657,44 @@ 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 + assert response.groups[0].classifier_cost == response.totals.classifier_cost == 0.4 + assert response.totals.spend - response.totals.classifier_cost == pytest.approx(9.6) + + @pytest.mark.asyncio + @pytest.mark.parametrize("recorded_turns", [0, 3, 10]) + async def test_classifier_subtotals_require_every_included_turn_to_be_recorded( + self, recorded_turns: int, monkeypatch: pytest.MonkeyPatch + ): + other: Final = self.ROW.model_copy( + update={ + "router_name": "other-auto", + "sessions": 1, + "turns": 10, + "spend": 2.0, + "saved_spend": -0.5, + "classifier_cost": recorded_turns * 0.02, + "classifier_cost_recorded_turns": recorded_turns, + } + ) + response: Final = await self._benchmarks( + monkeypatch, rows=[self.ROW.model_dump(), other.model_dump()], model_list=[] + ) + wire: Final = response.model_dump() + assert wire["groups"][0]["classifier_cost"] == 0.4 + assert wire["groups"][1]["classifier_cost"] == (pytest.approx(0.2) if recorded_turns == 10 else None) + assert wire["totals"]["classifier_cost"] == (pytest.approx(0.6) if recorded_turns == 10 else None) + assert response.totals.turns == 50 + assert response.totals.spend == 12.0 + assert response.totals.saved_spend == 29.5 + assert response.totals.baseline_spend == 41.5 + assert response.totals.saved_pct == 71.1 + assert response.totals.saved_per_session == 5.9 @pytest.mark.asyncio @pytest.mark.parametrize( @@ -720,6 +758,7 @@ class TestAutoRouterBenchmarks: assert (idle.cache.hit_rate_pct, idle.cache.coverage_pct) == (0.0, 0.0) assert idle.cache.same_model.turns == idle.cache.return_to_tier.hits == 0 assert idle.tier_turns == {} + assert idle.classifier_cost == 0.0 @pytest.mark.asyncio @pytest.mark.parametrize( @@ -796,7 +835,6 @@ class TestAutoRouterBenchmarks: from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock - from litellm.proxy.management_endpoints.auto_router_endpoints import ( get_shadow_eval_job, list_shadow_eval_jobs, @@ -1131,7 +1169,9 @@ async def test_start_shadow_eval_writes_one_leg_per_key_in_one_statement(monkeyp assert ( len( { - frozenset((k, tuple(v) if isinstance(v, list) else v) for k, v in row.items() if k not in ("target_id", "id")) + frozenset( + (k, tuple(v) if isinstance(v, list) else v) for k, v in row.items() if k not in ("target_id", "id") + ) for row in rows } ) @@ -1258,8 +1298,8 @@ async def test_start_shadow_eval_accepts_an_sdk_judge_when_anthropic_secret_look monkeypatch: pytest.MonkeyPatch, ) -> None: import litellm - from litellm.integrations.custom_secret_manager import CustomSecretManager import litellm.proxy.proxy_server as proxy_server + from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem class AnthropicSecretManager(CustomSecretManager): @@ -1821,9 +1861,10 @@ async def test_list_shadow_eval_jobs_rejects_a_lone_filter_half(monkeypatch: pyt @pytest.mark.asyncio async def test_start_shadow_eval_concurrent_unique_violation_is_a_409(monkeypatch: pytest.MonkeyPatch): - import litellm.proxy.proxy_server as proxy_server from prisma.errors import UniqueViolationError + import litellm.proxy.proxy_server as proxy_server + _configure_anthropic_sdk_judge(monkeypatch) prisma = _shadow_prisma() prisma.db.litellm_shadowevaljob.create_many = AsyncMock( diff --git a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py index 81226981089..0ecc4f8d7cb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py +++ b/tests/test_litellm/proxy/management_endpoints/test_credential_migration.py @@ -8,6 +8,7 @@ proof-of-fix (real proxy + DB) is performed separately on the repro server. import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -457,6 +458,65 @@ async def test_scan_covered_tables_classifies_legacy_and_v2(salt_key, monkeypatc assert by_loc["credentials"].legacy == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("column", ("static_headers", "env")) +@pytest.mark.parametrize("algorithm", ("xsalsa20-poly1305", "aes-256-gcm")) +@pytest.mark.parametrize("as_json", (False, True)) +@pytest.mark.parametrize( + "case", ("legacy", "encrypted", "wrong-key", "corrupt", "invalid-shape", "invalid-scalar", "empty", "null") +) +async def test_check_classifies_mcp_secret_maps( + salt_key: str, + monkeypatch: pytest.MonkeyPatch, + column: str, + algorithm: str, + as_json: bool, + case: str, +) -> None: + monkeypatch.setattr(proxy_server, "general_settings", {"encryption_algorithm": algorithm}) + plaintext: Final = {"Authorization": "v2:gcm:operator-text", "CUSTOM": "litellm_enc::literal\n café "} + ciphertext: Final = encrypt_value_helper(json.dumps(plaintext)) + cases: Final[dict[str, object]] = { + "legacy": plaintext, + "encrypted": ciphertext, + "wrong-key": encrypt_value_helper(json.dumps(plaintext), new_encryption_key="different-map-salt"), + "corrupt": ciphertext[:-4] + "AAAA", + "invalid-shape": encrypt_value_helper(json.dumps({"Authorization": 42})), + "invalid-scalar": "null", + "empty": {}, + "null": None, + } + value: Final = json.dumps(cases[case]) if as_json and case != "null" else cases[case] + row: Final = SimpleNamespace(**{column: value}) + client: Final = MagicMock() + _empty_covered_tables(client) + client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[row]) + client.db.litellm_mcpservertable.update = AsyncMock() + client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) + client.db.litellm_config.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "general_settings", {}) + + report: Final = await cm.check_encryption(client) + expected_legacy: Final = int(case == "legacy" or (case == "encrypted" and algorithm == "xsalsa20-poly1305")) + expected_v2: Final = int(case == "encrypted" and algorithm == "aes-256-gcm") + expected_invalid: Final = int(case in ("wrong-key", "corrupt", "invalid-shape", "invalid-scalar")) + + assert report.as_dict()["locations"]["mcp_server"] == { + "scanned": int(case not in ("empty", "null")), + "migrated": 0, + "already_v2": expected_v2, + "plaintext": 0, + "undecryptable": expected_invalid, + "legacy": expected_legacy, + } + assert report.residual_legacy == expected_legacy + assert report.total_undecryptable == expected_invalid + assert getattr(row, column) == value + client.db.litellm_mcpservertable.update.assert_not_awaited() + + @pytest.mark.asyncio async def test_check_counts_covered_table_residual(salt_key, monkeypatch): """check_encryption now scans the rotation-covered tables (model table here), 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/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index acb45038df0..d9969dd1dc9 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -5206,3 +5206,37 @@ class TestAzureRouterModelStreamingKeepalive: assert result.headers["x-upstream"] == "kept" assert chunks == [b"data: hello\n\n"] + + +@pytest.mark.asyncio +async def test_bedrock_count_tokens_error_forwards_provider_headers(): + """The count tokens route converts BedrockError into an HTTPException, and dropping the + headers there loses x-amzn-RequestId after the handler went to the trouble of keeping it.""" + from fastapi import HTTPException + + from litellm.llms.bedrock.common_utils import BedrockError + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + handle_bedrock_count_tokens, + ) + + failure = BedrockError( + status_code=500, + message="Amazon Bedrock is unable to process your request.", + headers={"x-amzn-RequestId": "req-count-tokens-500"}, + ) + + with patch( # test-quality-ok: the route's BedrockError branch is only reachable when the handler raises + "litellm.llms.bedrock.count_tokens.handler.BedrockCountTokensHandler.handle_count_tokens_request", + new=AsyncMock(side_effect=failure), + ): + with pytest.raises(HTTPException) as exc_info: + await handle_bedrock_count_tokens( + endpoint="v1/messages/count_tokens", + request=MagicMock(), + fastapi_response=MagicMock(), + user_api_key_dict=MagicMock(), + request_body={"model": "anthropic.claude-haiku-4-5-20251001-v1:0"}, + ) + + assert exc_info.value.status_code == 500 + assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-count-tokens-500" diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 4fcb7d22588..e95d001bb12 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -4,20 +4,26 @@ Tests for the pipeline executor. Uses mock guardrails to validate pipeline execution without external services. """ +import copy +from typing import Literal from unittest.mock import MagicMock import pytest import litellm +from litellm.caching.dual_cache import DualCache from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( CustomCodeGuardrail, ) from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) +from litellm.types.utils import CallTypesLiteral try: from fastapi.exceptions import HTTPException @@ -158,11 +164,146 @@ class ContentCheckGuardrail(CustomGuardrail): return None +class RecordingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str, scan_raw_request: bool = False, block: bool = True): + super().__init__( + guardrail_name=guardrail_name, + event_hook="pre_call", + default_on=True, + scan_raw_request=scan_raw_request, + ) + self.block = block + + def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool: + return True + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"detected": ["aws_access_key"]}, + request_data=data, + guardrail_status="guardrail_intervened" if self.block else "success", + ) + if self.block: + raise HTTPException(status_code=400, detail="Content policy violation") + return copy.deepcopy(data) + + # ───────────────────────────────────────────────────────────────────────────── # Tests # ───────────────────────────────────────────────────────────────────────────── +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +@pytest.mark.parametrize("scan_raw_request", [False, True]) +@pytest.mark.parametrize("on_fail", ["block", "modify_response"]) +async def test_terminal_block_carries_guardrail_information_to_request( + monkeypatch: pytest.MonkeyPatch, scan_raw_request: bool, on_fail: Literal["block", "modify_response"] +): + """ + Spend logging and the Guardrails Monitor read standard_logging_guardrail_information + off the caller's request dict. A blocking step records it on the executor's + working copy (or the raw-request snapshot), so the terminal result must carry it + back onto the request or the block is never counted. + """ + guard = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=scan_raw_request) + monkeypatch.setattr(litellm, "callbacks", [guard]) + data = { + "messages": [{"role": "user", "content": "key AKIAIOSFODNN7EXAMPLE"}], + "metadata": {"user_api_key_hash": "abc"}, + } + + result = await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="credentials-api-keys", on_fail=on_fail, on_pass="next")], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="baseline-pii-protection", + raw_request_snapshot={"messages": data["messages"], "metadata": {"user_api_key_hash": "abc"}}, + ) + + assert result.terminal_action == on_fail + recorded = data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["credentials-api-keys"] + assert recorded[0]["guardrail_status"] == "guardrail_intervened" + assert data["metadata"]["user_api_key_hash"] == "abc" + assert "guardrails" not in data["metadata"] + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_terminal_block_merges_guardrail_information_without_duplicates(monkeypatch: pytest.MonkeyPatch): + """A pass_data step that returns a rewritten copy of the request, and a scan_raw_request step + that evaluates a deep copy taken before the pipeline ran, both leave earlier entries in two + dicts at once. Those must be carried back once while every step's own entry is kept.""" + first = RecordingGuardrail(guardrail_name="pii-scan", block=False) + second = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=True) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"} + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + data["metadata"]["standard_logging_guardrail_information"] = [earlier] + + result = await PipelineExecutor.execute_steps( + steps=[ + PipelineStep(guardrail="pii-scan", on_fail="block", on_pass="next", pass_data=True), + PipelineStep(guardrail="credentials-api-keys", on_fail="block", on_pass="next"), + ], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="baseline-pii-protection", + raw_request_snapshot={ + "messages": data["messages"], + "metadata": {"standard_logging_guardrail_information": [dict(earlier)]}, + }, + ) + + assert result.terminal_action == "block" + recorded = data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "pii-scan", "credentials-api-keys"] + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_repeated_scan_raw_request_step_is_counted_once_per_evaluation(monkeypatch: pytest.MonkeyPatch): + """Running the same raw-scan guardrail twice yields two identical entries; both must reach the caller, + while the entries the raw snapshot already held before the pipeline ran are not copied again.""" + guard = RecordingGuardrail(guardrail_name="credentials-raw", scan_raw_request=True, block=False) + monkeypatch.setattr(litellm, "callbacks", [guard]) + earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"} + data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}} + data["metadata"]["standard_logging_guardrail_information"] = [earlier] + + result = await PipelineExecutor.execute_steps( + steps=[ + PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"), + PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"), + ], + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="raw-scan-policy", + raw_request_snapshot={ + "messages": data["messages"], + "metadata": {"standard_logging_guardrail_information": [dict(earlier)]}, + }, + ) + + assert result.terminal_action == "allow" + assert result.modified_data is not None + recorded = result.modified_data["metadata"]["standard_logging_guardrail_information"] + assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "credentials-raw", "credentials-raw"] + + @pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") @pytest.mark.asyncio async def test_escalation_step1_fails_step2_blocks(monkeypatch): diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a518c9289ae..dafe686c4f2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -22,6 +22,7 @@ import inspect import json import logging import os +import subprocess from collections.abc import Awaitable, Callable from typing import List, Optional, Union from unittest.mock import AsyncMock, MagicMock, patch @@ -749,6 +750,30 @@ async def test_proxy_startup_event_invalid_missing_app_arg_raises(): pass +@pytest.mark.asyncio +async def test_proxy_startup_event_prunes_dead_workers_live_gauges(tmp_path): + """With PROMETHEUS_MULTIPROC_DIR set, a booting worker drops the live-gauge files of pids that no longer + exist, so a crashed worker's in-flight samples leave the aggregate as soon as its replacement starts.""" + exited = subprocess.Popen(["true"]) + assert exited.wait(timeout=30) == 0 + stale = tmp_path / f"gauge_livesum_{exited.pid}.db" + stale.touch() + counter = tmp_path / f"counter_{exited.pid}.db" + counter.touch() + + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} + clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path) + with patch.dict(os.environ, clean_env, clear=True): + try: + async with proxy_startup_event(app=None): + pass + except Exception: + pass + + assert not stale.exists() + assert counter.exists() + + def test_otel_global_provider_published_after_callback_init(): """The OTel V2 global-provider publish must run after callback initialization in ``proxy_startup_event``. @@ -821,6 +846,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/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 2babfe432f3..ae4ec4086bb 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -9,6 +9,7 @@ Pins covered: from __future__ import annotations import json +import logging import os import re from types import SimpleNamespace @@ -1633,6 +1634,32 @@ async def test_ProxyConfig_load_config_minimal_yaml(tmp_path, monkeypatch): } +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", ["true", "false", "null", "'true'", None]) +async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monkeypatch, caplog, setting): + config_file = tmp_path / "budget.yaml" + flag = f" disable_budget_reservation: {setting}\n" if setting is not None else "" + config_file.write_text( + "model_list: []\nlitellm_settings: {}\ngeneral_settings:\n" + " master_key: null\n" + flag + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr("litellm.constants.budget_reservation_disabled_info_emitted", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + config = ProxyConfig() + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + for _ in range(3): + await config.load_config(router=None, config_file_path=str(config_file)) + + records = [ + record for record in caplog.records + if "disable_budget_reservation is enabled" in record.message + ] + assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else []) + + @pytest.mark.asyncio async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch): """Regression: router_settings.plugins dotted-path strings must be resolved to @@ -2912,6 +2939,129 @@ async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch fake_router.update_settings.assert_called_once_with(routing_strategy="latency-based-routing") +def _stub_add_deployment_collaborators( + monkeypatch: pytest.MonkeyPatch, pc: ProxyConfig, fake_prisma: MagicMock +) -> None: + from litellm.proxy import proxy_server + + fake_router = MagicMock() + fake_router.get_model_list = MagicMock(return_value=[]) + + async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]: + return {} + + monkeypatch.setattr(litellm, "credential_list", []) + monkeypatch.setattr(pc, "get_config", fake_get_config) + monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock()) + monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock()) + monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None)) + monkeypatch.setattr(proxy_server, "llm_router", fake_router) + monkeypatch.setattr(proxy_server, "master_key", "sk-master") + monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server, "proxy_config", pc) + monkeypatch.delenv("LITELLM_SALT_KEY", raising=False) + + +def _encrypted_credential_row(credential_name: str, api_key: str) -> dict[str, object]: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + return { + "credential_name": credential_name, + "credential_values": {"api_key": encrypt_value_helper(api_key, new_encryption_key="sk-master")}, + "credential_info": {"custom_llm_provider": "openai"}, + } + + +def _fake_prisma_with_encrypted_credential(credential_name: str, api_key: str) -> MagicMock: + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock( + return_value=[_encrypted_credential_row(credential_name, api_key)] + ) + return fake_prisma + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconciling_models( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy import proxy_server + from litellm.utils import load_credentials_from_list + + pc = ProxyConfig() + fake_prisma = MagicMock() + fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {}) + installed = MagicMock() + + async def read_models_while_a_credential_lands(prisma_client: object) -> list[MagicMock]: + fake_prisma.db.litellm_credentialstable.find_many.return_value = [ + _encrypted_credential_row("openai-cred", "sk-from-db") + ] + return [MagicMock()] + + async def install_models(new_models: object, proxy_logging_obj: object) -> None: + installed(credential=CredentialAccessor.get_credential_values("openai-cred")) + + monkeypatch.setattr(pc, "_get_models_from_db", read_models_while_a_credential_lands) + monkeypatch.setattr(pc, "_update_llm_router", install_models) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + installed.assert_called_once_with(credential={"api_key": "sk-from-db"}) + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} + request_kwargs = {"litellm_credential_name": "openai-cred"} + load_credentials_from_list(request_kwargs) + assert request_kwargs == {"litellm_credential_name": "openai-cred", "api_key": "sk-from-db"} + + +@pytest.mark.asyncio +async def test_ProxyConfig_add_deployment_loads_db_credentials_even_when_models_are_not_db_objects( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy import proxy_server + + pc = ProxyConfig() + fake_prisma = _fake_prisma_with_encrypted_credential("openai-cred", "sk-from-db") + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["mcp"]}) + models_fetch = AsyncMock(return_value=[]) + monkeypatch.setattr(pc, "_get_models_from_db", models_fetch) + + await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock()) + + models_fetch.assert_not_awaited() + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"} + + +@pytest.mark.asyncio +async def test_ProxyConfig_get_credentials_reads_from_writer_not_replica(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + pc = ProxyConfig() + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + writer_inner.litellm_credentialstable.find_many = AsyncMock( + return_value=[_encrypted_credential_row("openai-cred", "sk-from-writer")] + ) + reader_inner.litellm_credentialstable.find_many = AsyncMock(return_value=[]) + fake_prisma = MagicMock() + fake_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + _stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma) + + await pc.get_credentials(prisma_client=fake_prisma) + + assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-writer"} + reader_inner.litellm_credentialstable.find_many.assert_not_awaited() + + # --------------------------------------------------------------------------- # ProxyConfig._add_general_settings_from_db_config # --------------------------------------------------------------------------- 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..e466edab131 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") @@ -1021,6 +1162,104 @@ def test_prompt_caching_prices_at_the_deployment_rate_not_the_public_one(): assert result.prompt_caching > at_public_rates.prompt_caching +@pytest.mark.parametrize( + "baseline_id, selected_id, selected_multiplier, billed_input, classifier_cost, expected", + [ + ("baseline", "selected", 0.1, None, 0.0, 0.0135), + ("baseline", "selected", 2.0, None, 0.0, -0.015), + ("baseline", "selected", 1.0, None, 0.0, 0.0), + ("baseline", "selected", 0.1, 0.004, 0.001, 0.01), + ("baseline", "baseline", 0.1, 0.004, 0.001, -0.001), + (None, "selected", 0.1, None, 0.0, 0.0), + ("baseline", None, 0.1, None, 0.0, 0.0), + (None, None, 0.1, None, 0.0, 0.0), + ("", "selected", 0.1, None, 0.0, 0.0), + ("baseline", "", 0.1, None, 0.0, 0.0), + ], +) +def test_autorouter_savings_distinguishes_priced_deployments( + baseline_id: str | None, + selected_id: str | None, + selected_multiplier: float, + billed_input: float | None, + classifier_cost: float, + expected: float, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": "anthropic/claude-opus-5", + "api_key": "test-key", + "input_cost_per_token": 1e-5 * multiplier, + "output_cost_per_token": 5e-5 * multiplier, + }, + "model_info": {"id": name}, + } + for name, multiplier in (("baseline", 1.0), ("selected", selected_multiplier)) + ] + ) + result: Final = compute_savings_spend( + model="claude-opus-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=False, + model_id=selected_id, + llm_router=lambda: router, + routing_decision={ + "savings_baseline_model": "anthropic/claude-opus-5", + "savings_baseline_deployment_id": baseline_id, + "conversation_continuing": False, + "classifier_cost": classifier_cost, + }, + usage_object={"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, + cost_breakdown=None if billed_input is None else {"input_cost": billed_input, "output_cost": 0.0}, + ) + assert result.autorouter == pytest.approx(expected) + + +@pytest.mark.parametrize("selected_model", ["azure/contract-deployment", "contract-deployment"]) +def test_autorouter_savings_recognizes_one_deployment_under_its_base_model(selected_model: str) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "contract", + "litellm_params": { + "model": "azure/contract-deployment", + "api_key": "test-key", + "api_base": "https://example.openai.azure.com", + "input_cost_per_token": 0.0001, + "output_cost_per_token": 0.0002, + "cache_read_input_token_cost": 0.00001, + }, + "model_info": {"id": "contract", "base_model": "azure/gpt-5.5"}, + } + ] + ) + result: Final = compute_savings_spend( + model=selected_model, + custom_llm_provider="azure", + compression_saved_tokens=0, + gateway_injected_cache=False, + model_id="contract", + llm_router=lambda: router, + routing_decision={ + "savings_baseline_model": "azure/gpt-5.5", + "savings_baseline_deployment_id": "contract", + "conversation_continuing": True, + }, + usage_object={ + "prompt_tokens": 21000, + "completion_tokens": 100, + "total_tokens": 21100, + "prompt_tokens_details": {"text_tokens": 1000, "cached_tokens": 0, "cache_creation_tokens": 20000}, + }, + cost_breakdown={"input_cost": 2.1, "output_cost": 0.02}, + ) + assert result.autorouter == 0.0 + + def test_a_recorded_baseline_deployment_prices_at_its_configured_rate(): """A hardest-tier deployment with a negotiated rate is what the traffic would really have cost; pricing its model publicly misstates the saving.""" diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 96d9b0c5a26..bfae42f64f1 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,143 @@ 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 + + +@pytest.mark.asyncio +async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status_error(): + """The httpx.HTTPStatusError branch dropped the headers its sibling branches forward. + + A Bedrock passthrough failure reaches this branch, so the request id was gone + before the client saw the response. + """ + import httpx + + from litellm.proxy._types import UserAPIKeyAuth + + request = httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse") + response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-passthrough-500"}, + content=b'{"message": "Amazon Bedrock is unable to process your request."}', + request=request, + ) + + processor = ProxyBaseLLMRequestProcessing(data={}) + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + with pytest.raises(HTTPException) as exc_info: + await processor._handle_llm_api_exception( + e=httpx.HTTPStatusError("boom", request=request, response=response), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=proxy_logging_obj, + ) + + assert exc_info.value.headers is not None + assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 7070617ce3e..74799072984 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -7272,7 +7272,12 @@ def _reserved_stamp_key(key_metadata: dict | None = None) -> UserAPIKeyAuth: ) -_PLANTED_STAMPS = {"attempted_fallbacks": 99, "original_model_group": "spoofed-group", "client_key": "client_value"} +_PLANTED_STAMPS = { + "attempted_fallbacks": 99, + "original_model_group": "spoofed-group", + "_client_output_ceiling": {"api_base": "https://attacker.example"}, + "client_key": "client_value", +} @pytest.mark.asyncio @@ -7301,6 +7306,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo assert "litellm_metadata" not in updated assert "attempted_fallbacks" not in updated["metadata"] assert "original_model_group" not in updated["metadata"] + assert "_client_output_ceiling" not in updated["metadata"] assert updated["metadata"]["client_key"] == "client_value" diff --git a/tests/test_litellm/proxy/test_prometheus_cleanup.py b/tests/test_litellm/proxy/test_prometheus_cleanup.py index 93b9b694c2c..6a1b95c51ff 100644 --- a/tests/test_litellm/proxy/test_prometheus_cleanup.py +++ b/tests/test_litellm/proxy/test_prometheus_cleanup.py @@ -6,13 +6,77 @@ ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir. from __future__ import annotations import os +import subprocess +import sys +from pathlib import Path +from typing import Final from unittest.mock import patch import pytest +from prometheus_client import CollectorRegistry, multiprocess -from litellm.proxy.prometheus_cleanup import mark_worker_exit, wipe_directory +from litellm.proxy.prometheus_cleanup import mark_dead_workers, mark_worker_exit, wipe_directory from litellm.proxy.proxy_cli import ProxyInitializationHelpers +_WORKER: Final = """ +import sys, time +from prometheus_client import Gauge +Gauge("litellm_in_flight", "", multiprocess_mode="livesum").set(float(sys.argv[1])) +print("ready", flush=True) +if sys.argv[2] == "stay": + time.sleep(120) +""" + + +def _spawn_worker(directory: Path, in_flight: str, lifetime: str) -> subprocess.Popen[str]: + env = {**os.environ, "PROMETHEUS_MULTIPROC_DIR": str(directory)} + worker = subprocess.Popen( + [sys.executable, "-c", _WORKER, in_flight, lifetime], env=env, stdout=subprocess.PIPE, text=True + ) + assert worker.stdout is not None and worker.stdout.readline() == "ready\n" + return worker + + +def _livesum(directory: Path) -> float: + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry, path=str(directory)) + value = registry.get_sample_value("litellm_in_flight") + return 0.0 if value is None else value + + +class TestMarkDeadWorkers: + def test_drops_live_gauges_of_exited_workers_and_keeps_running_ones(self, tmp_path: Path) -> None: + """A worker that died mid-request leaves its livesum file behind; the replacement worker's startup prune + must remove exactly that file so the aggregate stops counting requests nobody is serving.""" + dead = _spawn_worker(tmp_path, "3", "exit") + assert dead.wait(timeout=30) == 0 + alive = _spawn_worker(tmp_path, "2", "stay") + try: + assert (tmp_path / f"gauge_livesum_{dead.pid}.db").exists() + assert _livesum(tmp_path) == 5.0 + + assert mark_dead_workers(str(tmp_path)) == (dead.pid,) + + assert not (tmp_path / f"gauge_livesum_{dead.pid}.db").exists() + assert (tmp_path / f"gauge_livesum_{alive.pid}.db").exists() + assert _livesum(tmp_path) == 2.0 + assert mark_dead_workers(str(tmp_path)) == () + finally: + alive.kill() + alive.wait(timeout=30) + + def test_leaves_counters_of_exited_workers_alone(self, tmp_path: Path) -> None: + (tmp_path / "counter_424242.db").touch() + (tmp_path / "histogram_424242.db").touch() + assert mark_dead_workers(str(tmp_path)) == () + assert sorted(p.name for p in tmp_path.glob("*.db")) == ["counter_424242.db", "histogram_424242.db"] + + def test_keeps_live_gauges_of_workers_it_may_not_signal(self, tmp_path: Path) -> None: + """Signal 0 to pid 1 raises PermissionError for an unprivileged proxy; that pid is alive, not dead.""" + (tmp_path / "gauge_livesum_1.db").touch() + assert mark_dead_workers(str(tmp_path)) == () + assert (tmp_path / "gauge_livesum_1.db").exists() + class TestWipeDirectory: def test_deletes_all_db_files(self, tmp_path): diff --git a/tests/test_litellm/proxy/test_prometheus_metrics_server.py b/tests/test_litellm/proxy/test_prometheus_metrics_server.py index fc1fa381fa4..e2f461e61a6 100644 --- a/tests/test_litellm/proxy/test_prometheus_metrics_server.py +++ b/tests/test_litellm/proxy/test_prometheus_metrics_server.py @@ -1,5 +1,5 @@ -"""The separate metrics server must aggregate PROMETHEUS_MULTIPROC_DIR, expose only /metrics, and follow its -parent's lifetime. +"""The separate metrics server must aggregate PROMETHEUS_MULTIPROC_DIR, expose /metrics plus a probe-friendly +/health, and follow its parent's lifetime. Everything here runs on loopback against a child of this test process; no LLM keys or external network. """ @@ -91,7 +91,10 @@ def test_metrics_app_aggregates_multiproc_dir_and_reports_pid(tmp_path: Path, mo assert metrics.headers[PID_HEADER] == str(os.getpid()) assert 'litellm_requests_metric_total{model="gpt-5"} 5.0' in metrics.text - assert client.get("/health").status_code == 404 + health: Final = client.get("/health") + assert health.status_code == 200 + assert health.json() == {"status": "healthy", "multiproc_dir": str(tmp_path)} + assert client.get("/docs").status_code == 404 empty: Final = TestClient(build_metrics_app(str(other_dir))).get("/metrics") assert empty.status_code == 200 diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 9f1321aec2c..22930a26974 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -637,14 +637,14 @@ def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overri assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == [] -def test_deployment_pre_call_target_stays_native_when_opted_out(): +def test_deployment_hook_target_stays_native_when_opted_out(): """Model-level guardrails resolve their target here rather than through ProxyLogging.""" - assert _KeepsNativeHooks()._deployment_pre_call_target() is not None + assert _KeepsNativeHooks()._deployment_hook_target() is not None opted_out = _KeepsNativeHooks() - assert opted_out._deployment_pre_call_target() is opted_out - assert _AppliesGuardrail()._deployment_pre_call_target() is not None + assert opted_out._deployment_hook_target() is opted_out + assert _AppliesGuardrail()._deployment_hook_target() is not None routed = _AppliesGuardrail() - assert routed._deployment_pre_call_target() is not routed + assert routed._deployment_hook_target() is not routed @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index b7bb58378d4..6a901c5e7de 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -822,16 +822,16 @@ def _mock_scheduled_proxy_config() -> MagicMock: @pytest.mark.asyncio -async def test_initialize_scheduled_jobs_credentials(monkeypatch): - """ - Test that get_credentials is only called when store_model_in_db is True - """ +async def test_initialize_scheduled_jobs_loads_credentials_only_through_add_deployment( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from apscheduler.schedulers.asyncio import AsyncIOScheduler + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging - # Mock dependencies mock_prisma_client = MagicMock() mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() @@ -841,25 +841,6 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): with ( patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch("litellm.proxy.proxy_server.store_model_in_db", False), - ): # set store_model_in_db to False - # Test when store_model_in_db is False - await ProxyStartupEvent.initialize_scheduled_background_jobs( - general_settings={}, - prisma_client=mock_prisma_client, - proxy_budget_rescheduler_min_time=1, - proxy_budget_rescheduler_max_time=2, - proxy_batch_write_at=5, - proxy_logging_obj=mock_proxy_logging, - ) - - # Verify get_credentials was not called - mock_proxy_config.get_credentials.assert_not_called() - - # Now test with store_model_in_db = True - with ( - patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), - patch("litellm.proxy.proxy_server.store_model_in_db", True), - patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True), ): await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings={}, @@ -870,12 +851,31 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): proxy_logging_obj=mock_proxy_logging, ) - # Verify get_credentials was called both directly and scheduled - assert mock_proxy_config.get_credentials.call_count == 1 # Direct call + mock_proxy_config.get_credentials.assert_not_called() + mock_proxy_config.add_deployment.assert_not_called() - # Verify a scheduled job was added for get_credentials - mock_scheduler_calls = [call[0] for call in mock_proxy_config.get_credentials.mock_calls] - assert len(mock_scheduler_calls) > 0 + scheduler = AsyncIOScheduler() + try: + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), + ): + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + assert scheduler.get_job("get_credentials_job") is None + assert scheduler.get_job("add_deployment_job") is not None + mock_proxy_config.get_credentials.assert_not_called() + assert mock_proxy_config.add_deployment.call_count == 1 + finally: + scheduler.shutdown(wait=False) @pytest.mark.asyncio @@ -924,7 +924,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat @pytest.mark.asyncio async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch): """ - The DB config-reload jobs (add_deployment, get_credentials) that keep multi-pod + The DB config-reload job (add_deployment) that keeps multi-pod deployments in sync must be scheduled at the configured proxy_config_reload_interval_seconds, not a hardcoded value. """ @@ -967,7 +967,7 @@ async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval( if "id" in job_call.kwargs } assert scheduled_seconds["add_deployment_job"] == configured_interval - assert scheduled_seconds["get_credentials_job"] == configured_interval + assert "get_credentials_job" not in scheduled_seconds @pytest.mark.asyncio @@ -1011,7 +1011,7 @@ async def test_initialize_scheduled_jobs_rejects_non_positive_config_reload_inte if "id" in job_call.kwargs } assert scheduled_seconds["add_deployment_job"] == 30 - assert scheduled_seconds["get_credentials_job"] == 30 + assert "get_credentials_job" not in scheduled_seconds @pytest.mark.asyncio @@ -3166,6 +3166,47 @@ async def test_custom_ui_sso_sign_in_handler_config_loading(): os.unlink(config_file_path) +@pytest.mark.asyncio +async def test_startup_initializes_string_callbacks_after_all_litellm_settings_load(tmp_path, monkeypatch): + from litellm.integrations.s3_v2 import S3Logger + from litellm.litellm_core_utils import litellm_logging + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import ProxyLogging + + config_file = tmp_path / "config.yaml" + config_file.write_text( + "model_list: []\n" + "litellm_settings:\n" + " success_callback:\n" + " - s3_v2\n" + " failure_callback:\n" + " - s3_v2\n" + " s3_callback_params:\n" + " s3_bucket_name: ordering-regression-bucket\n" + " s3_region_name: us-west-2\n" + ) + + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "s3_callback_params", None) + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) + + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + ProxyLogging(user_api_key_cache=MagicMock())._init_litellm_callbacks(llm_router=None) + + success_loggers = [cb for cb in litellm._async_success_callback if isinstance(cb, S3Logger)] + failure_loggers = [cb for cb in litellm._async_failure_callback if isinstance(cb, S3Logger)] + assert len(success_loggers) == 1 + assert len(failure_loggers) == 1 + assert success_loggers[0].s3_bucket_name == "ordering-regression-bucket" + assert success_loggers[0].s3_region_name == "us-west-2" + assert "s3_v2" not in litellm.success_callback + assert "s3_v2" not in litellm.failure_callback + + @pytest.mark.asyncio async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeypatch): """ @@ -7446,10 +7487,8 @@ async def test_store_model_in_db_db_override_when_config_false(): # store_model_in_db should now be True (overridden by DB) assert ps.store_model_in_db is True - # add_deployment and get_credentials should have been called - # since store_model_in_db is now True assert mock_proxy_config.add_deployment.call_count == 1 - assert mock_proxy_config.get_credentials.call_count == 1 + mock_proxy_config.get_credentials.assert_not_called() @pytest.mark.asyncio 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/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index ab4c5185057..2c1845f7b92 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -1,6 +1,13 @@ +import json +import sys +import types + import pytest +import respx +from httpx import Response from unittest.mock import AsyncMock, patch +import litellm from litellm.types.utils import ModelResponse from litellm.responses.mcp import chat_completions_handler @@ -1344,3 +1351,39 @@ async def test_acompletion_with_mcp_streaming_drains_inner_stream_after_exhausti assert len(all_chunks) == 3 assert initial_stream.drained_after_exhaustion is True + + +@pytest.mark.asyncio +@respx.mock +async def test_acompletion_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + + zapier_tool = {"type": "mcp", "server_label": "zapier", "server_url": "https://mcp.zapier.com/api/mcp/mcp"} + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + provider = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-zapier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4.1", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + + result = await acompletion_with_mcp( + model="openai/gpt-4.1", + messages=[{"role": "user", "content": "hello"}], + tools=[zapier_tool], + api_key="sk-test", + acompletion=True, + ) + + assert isinstance(result, ModelResponse) + assert result.id == "chatcmpl-zapier" + assert json.loads(provider.calls.last.request.content)["tools"] == [zapier_tool] diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py index 80151d0cba8..9745a0af970 100644 --- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException +from openai.types.responses.tool_param import Mcp import importlib from litellm.proxy._experimental.mcp_server.faults.list_outcomes import AggregateToolListing @@ -724,6 +725,178 @@ def test_extract_tool_call_details_still_prefers_openai_arguments(): assert arguments == '{"city": "Paris"}' +def _registered( + server_id: str, + name: str, + alias: str | None = None, + server_name: str | None = None, + access_groups: list[str] | None = None, +): + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + return MCPServer( + server_id=server_id, + name=name, + alias=alias, + server_name=server_name, + transport=MCPTransport.http, + access_groups=access_groups, + ) + + +async def _no_toolset(_: str) -> bool: + return False + + +ZAPIER_TOOL: Mcp = { + "type": "mcp", + "server_label": "zapier", + "server_url": "https://mcp.zapier.com/api/mcp/mcp", + "require_approval": "never", +} +EXPLICIT_GATEWAY_TOOL = {"type": "mcp", "server_label": "github", "server_url": "litellm_proxy/mcp/github"} +FUNCTION_TOOL = {"type": "function", "name": "get_weather", "parameters": {}} + + +@pytest.mark.asyncio +async def test_gateway_served_names_matches_alias_server_name_name_access_group_and_toolset(): + from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + + servers = ( + _registered("id-1", "github-name", alias="github", server_name="github-server", access_groups=["prod-group"]), + _registered("id-2", "deepwiki"), + ) + + async def toolset_exists(name: str) -> bool: + return name == "my-toolset" + + served = await _gateway_served_names( + {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset", "mcp", "nope"}, + servers=lambda: servers, + toolset_exists=toolset_exists, + ) + + assert served == {"github", "github-server", "github-name", "deepwiki", "prod-group", "my-toolset"} + + +@pytest.mark.asyncio +async def test_gateway_served_names_matches_server_id_short_prefix_and_alias_case_like_the_gateway(): + from litellm.proxy._experimental.mcp_server.utils import compute_short_server_prefix + from litellm.responses.mcp.litellm_proxy_mcp_handler import _gateway_served_names + + server_id = "0b9ae4ca-1bd2-4faa-b183-7dd812597e3b" + short_prefix = compute_short_server_prefix(server_id) + servers = (_registered(server_id, "github-name", alias="github", access_groups=["prod-group"]),) + + served = await _gateway_served_names( + {server_id, short_prefix, "GitHub", "PROD-GROUP", "nope"}, servers=lambda: servers, toolset_exists=_no_toolset + ) + + assert served == {server_id, short_prefix, "GitHub"} + + +@pytest.mark.asyncio +async def test_routes_through_gateway_flags_explicit_and_served_tools_only(): + served_tool = {"type": "mcp", "server_label": "github", "server_url": "http://localhost:4000/mcp/github"} + + async def served_names(names): + assert names == {"github", "mcp"} + return frozenset({"github"}) + + flags = await LiteLLM_Proxy_MCP_Handler.routes_through_gateway( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, served_tool, FUNCTION_TOOL], served_names=served_names + ) + + assert flags == (False, True, True, False) + + +@pytest.mark.asyncio +async def test_split_mcp_tools_leaves_external_mcp_path_urls_for_the_provider(): + + async def served_names(names): + assert names == {"mcp"} + return frozenset() + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [ZAPIER_TOOL, EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [ZAPIER_TOOL, FUNCTION_TOOL] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_repoints_served_proxy_urls_at_the_gateway(): + served_tool = { + "type": "mcp", + "server_label": "toolset", + "server_url": "http://localhost:4000/mcp/my-toolset", + "require_approval": "never", + "allowed_tools": ["get_me"], + } + unserved_tool = {"type": "mcp", "server_label": "typo", "server_url": "http://localhost:4000/mcp/githb"} + + async def served_names(names): + return frozenset({"my-toolset"}) + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [served_tool, unserved_tool], served_names=served_names + ) + + assert gateway_tools == [{**served_tool, "server_url": "litellm_proxy/mcp/my-toolset"}] + assert other_tools == [unserved_tool] + + +@pytest.mark.asyncio +async def test_split_mcp_tools_skips_resolution_when_nothing_points_at_the_proxy(): + async def served_names(names): + raise AssertionError("no lookup expected") + + gateway_tools, other_tools = await LiteLLM_Proxy_MCP_Handler._split_mcp_tools( + [EXPLICIT_GATEWAY_TOOL, FUNCTION_TOOL], served_names=served_names + ) + + assert gateway_tools == [EXPLICIT_GATEWAY_TOOL] + assert other_tools == [FUNCTION_TOOL] + + +def test_should_use_gateway_still_triggers_on_http_mcp_path(): + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([ZAPIER_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([EXPLICIT_GATEWAY_TOOL]) is True + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway([FUNCTION_TOOL]) is False + assert LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway(None) is False + + +@pytest.mark.asyncio +async def test_aresponses_api_with_mcp_forwards_unserved_external_mcp_tool_to_the_provider(monkeypatch): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.responses import main as responses_main + from litellm.types.llms.openai import ResponsesAPIResponse + + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", types.SimpleNamespace(prisma_client=None)) + monkeypatch.setattr(global_mcp_server_manager, "get_registry", lambda: {}) + provider_tools: list[object] = [] + + def fake_provider(**kwargs: object) -> object: + request_params = cast(dict[str, object], kwargs["response_api_optional_request_params"]) + provider_tools.append(request_params.get("tools")) + + async def respond() -> ResponsesAPIResponse: + return ResponsesAPIResponse(id="resp_zapier", created_at=0, output=[]) + + return respond() + + monkeypatch.setattr(responses_main.base_llm_http_handler, "response_api_handler", fake_provider) + + response = await responses_main.aresponses_api_with_mcp( + input="Reply with the single word ok.", model="openai/gpt-4.1", tools=[ZAPIER_TOOL] + ) + + assert isinstance(response, ResponsesAPIResponse) + assert provider_tools == [[ZAPIER_TOOL]] + + def _response_with_reasoning_and_tool_call() -> Any: """A first-turn response as a reasoning model returns it: reasoning item, then a function call.""" return ResponsesAPIResponse( diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index e8333214ea8..fe3c4a0640d 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -17,6 +17,9 @@ from litellm.llms.chatgpt.responses.transformation import ChatGPTResponsesAPICon from litellm.llms.databricks.responses.transformation import ( DatabricksResponsesAPIConfig, ) +from litellm.llms.fireworks_ai.responses.transformation import ( + FireworksAIResponsesAPIConfig, +) from litellm.llms.github_copilot.responses.transformation import ( GithubCopilotResponsesAPIConfig, ) @@ -102,6 +105,12 @@ class TestResponsesAPIWebSocketSupport: def test_openai_model_in_websocket_url_default(self): assert OpenAIResponsesAPIConfig().model_in_websocket_url() is True + def test_fireworks_ai_uses_managed_websocket(self): + """Fireworks AI should use managed websocket handler""" + assert ( + FireworksAIResponsesAPIConfig().supports_native_websocket() is False + ), "Fireworks AI should use managed websocket handler" + def test_xai_uses_managed_websocket(self): """XAI should use managed websocket handler""" config = XAIResponsesAPIConfig() diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index ebfb631f93b..565e77b0ae0 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -7,7 +7,8 @@ Tests the rule-based complexity scoring and tier assignment logic. import asyncio import logging import sys -from typing import Dict, List +import time +from typing import Dict, Final, List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -15,6 +16,7 @@ from pydantic import ValidationError import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.auto_router_model_naming import ( CUSTOMIZATION_CAPABILITY, GATED_AUTO_ROUTER_CAPABILITIES, @@ -23,7 +25,12 @@ from litellm.router_utils.auto_router_model_naming import ( ) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import ( + OUTPUT_TOKEN_CEILING_PARAMS, + RETURN_RAW_MODEL_NAME_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) +from litellm.router import as_output_cap from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -47,6 +54,7 @@ from litellm.router_strategy.complexity_router.config import ( ClassifierLLMConfig, ComplexityRouterConfig, ComplexityTier, + custom_pattern_work, ) from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, @@ -756,6 +764,289 @@ class TestCustomTechnicalKeywords: assert custom_score > baseline_score +class TestCustomDimensions: + @pytest.mark.parametrize( + "matchers,prompt", + [ + pytest.param( + {"keywords": ["orbitmesh", "fluxgate"]}, + "Connect ORBITMESH and fluxgate for the requested change", + id="keywords", + ), + pytest.param( + {"patterns": [r"\bCREATE\s{1,4}TABLE\b", r"\bALTER\s{1,4}TABLE\b"]}, + "create table widgets (id integer); ALTER TABLE widgets ADD label text;", + id="regex", + ), + ], + ) + def test_custom_dimension_changes_only_matching_requests( + self, mock_router_instance: MagicMock, matchers: dict[str, object], prompt: str + ) -> None: + baseline: Final = ComplexityRouter("test-router", mock_router_instance) + configured: Final = ComplexityRouter( + "test-router", + mock_router_instance, + {"custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, **matchers}]}, + ) + baseline_tier, baseline_score, baseline_signals = baseline.classify(prompt) + tier, score, signals = configured.classify(prompt) + assert baseline_tier == ComplexityTier.SIMPLE + assert tier != ComplexityTier.SIMPLE + assert score == pytest.approx(baseline_score + 0.7) + assert signals == [*baseline_signals, "custom (internalFrameworks)"] + plain: Final = "Hello!" + assert configured.classify(plain) == baseline.classify(plain) + assert configured.classify(plain)[0] == ComplexityTier.SIMPLE + + @pytest.mark.parametrize( + "dimension_overrides,config_overrides", + [ + pytest.param({"keywords": []}, {}, id="missing-matchers"), + pytest.param({"keywords": [" "]}, {}, id="blank-keyword"), + pytest.param({"patterns": ["\t"]}, {}, id="blank-pattern"), + pytest.param({"patterns": ["("]}, {}, id="invalid-regex"), + pytest.param({"patterns": [r"a*b"]}, {}, id="unbounded-star"), + pytest.param({"patterns": [r"a{2,}b"]}, {}, id="unbounded-brace"), + pytest.param({"patterns": [r"a{0,65}b"]}, {}, id="repeat-over-64"), + pytest.param({"patterns": [r"(a{0,8}){0,8}b"]}, {}, id="nested-repeat"), + pytest.param({"patterns": [r"(a|aa){0,12}b"]}, {}, id="alternation-in-repeat"), + pytest.param({"patterns": [r"(?:ab){0,64}c"]}, {}, id="group-repeat"), + pytest.param({"patterns": ["a?" * 9 + "b"]}, {}, id="pattern-work-over-budget"), + pytest.param({"patterns": ["(?:a|aa)" * 9 + "z"]}, {}, id="ambiguous-alternation-chain"), + pytest.param({"patterns": ["a?" * 8 + "a{64}" * 10 + "z"]}, {}, id="cheap-prefix-expensive-tail"), + pytest.param({"patterns": [r"(a)\1"]}, {}, id="backreference"), + pytest.param({"patterns": [r"(?=x)y"]}, {}, id="lookahead"), + pytest.param({"patterns": [r"(?>ab)"]}, {}, id="atomic-group"), + pytest.param({"patterns": [r"a*+b"]}, {}, id="possessive"), + pytest.param({"name": "CODEPRESENCE"}, {"dimension_weights": {"tokenCount": 0.1}}, id="reserved-name"), + pytest.param({}, {"dimension_weights": {"INTERNALFRAMEWORKS": 0.7}}, id="weight-in-map"), + pytest.param({"weight": 0}, {}, id="zero-weight"), + pytest.param({"weight": 1.1}, {}, id="excess-weight"), + pytest.param({"weight": float("nan")}, {}, id="nan-weight"), + pytest.param({"weight": float("inf")}, {}, id="infinite-weight"), + pytest.param({"name": "bad-name"}, {}, id="invalid-name"), + pytest.param({"name": "x" * 65}, {}, id="long-name"), + pytest.param({"keywords": [""]}, {}, id="empty-matcher"), + pytest.param({"keywords": ["x" * 257]}, {}, id="long-matcher"), + pytest.param({"keywords": ["x"] * 32, "patterns": ["y"]}, {}, id="combined-matcher-count"), + pytest.param({"keywords": ["x" * 256] * 17}, {}, id="matcher-character-budget"), + pytest.param({"unknown": True}, {}, id="extra-field"), + pytest.param({"scoring_mode": "graded"}, {}, id="unknown-scoring-mode"), + pytest.param({"scoring_mode": None}, {}, id="null-scoring-mode"), + ], + ) + def test_custom_dimension_invalid_configuration_rejected( + self, dimension_overrides: dict[str, object], config_overrides: dict[str, object] + ) -> None: + with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"): + ComplexityRouterConfig.model_validate( + { + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.7, + "keywords": ["orbitmesh"], + **dimension_overrides, + } + ], + **config_overrides, + } + ) + + @pytest.mark.parametrize( + "names", + [ + pytest.param(("internalFrameworks", "INTERNALFRAMEWORKS"), id="duplicate-casefolded-name"), + pytest.param(tuple(f"dimension{i}" for i in range(17)), id="dimension-count"), + ], + ) + def test_custom_dimension_names_and_count_are_bounded(self, names: tuple[str, ...]) -> None: + with pytest.raises(ValidationError, match=r"custom_dimensions|custom dimension"): + ComplexityRouterConfig.model_validate( + {"custom_dimensions": [{"name": name, "weight": 0.7, "keywords": ["orbitmesh"]} for name in names]} + ) + + @pytest.mark.parametrize("classifier_type", ("heuristic_v2", "llm", "custom")) + def test_custom_dimensions_reject_classifiers_outside_the_tuning_gate(self, classifier_type: str) -> None: + classifier_config: Final = ( + {"classifier_plugin": _FixedTierClassifier("SIMPLE")} + if classifier_type == "custom" + else {"classifier_llm_config": {"model": "judge"}} + if classifier_type == "llm" + else {} + ) + with pytest.raises(ValidationError, match="custom_dimensions requires classifier_type"): + ComplexityRouterConfig.model_validate( + { + "classifier_type": classifier_type, + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}], + **classifier_config, + } + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("scoring_mode", ("binary", "match_count")) + @pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh", "orbitmesh fluxgate")) + async def test_custom_dimensions_public_hook_scores_only_current_ask( + self, mock_router_instance: MagicMock, current_ask: str, scoring_mode: str + ) -> None: + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + { + "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}, + "dimension_weights": {}, + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.8, + "keywords": ["orbitmesh", "fluxgate"], + "scoring_mode": scoring_mode, + } + ], + }, + ) + result: Final = await router.async_pre_routing_hook( + model="test-router", + request_kwargs={}, + messages=[ + {"role": "system", "content": "orbitmesh fluxgate"}, + {"role": "user", "content": "orbitmesh fluxgate"}, + {"role": "assistant", "content": "orbitmesh fluxgate is ready"}, + {"role": "user", "content": current_ask}, + {"role": "tool", "tool_call_id": "previous", "content": "orbitmesh fluxgate"}, + ], + ) + assert result is not None + assert result.routing_decision is not None + expected_score: Final = ( + 0.0 + if current_ask == "Hello!" + else 0.4 + if scoring_mode == "match_count" and current_ask == "orbitmesh" + else 0.8 + ) + assert result.routing_decision["score"] == expected_score + assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (expected_score > 0) + assert result.model == ("cheap" if expected_score == 0 else "strong" if expected_score == 0.4 else "top") + assert "orbitmesh" not in " ".join(result.routing_decision["signals"]) + + @pytest.mark.parametrize("scoring_mode", ("binary", "match_count")) + def test_custom_patterns_scan_only_the_first_2048_characters( + self, mock_router_instance: MagicMock, scoring_mode: str + ) -> None: + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + { + "custom_dimensions": [ + { + "name": "late", + "weight": 0.7, + "patterns": [r"zzz{1,3}", r"yyy{1,3}"], + "scoring_mode": scoring_mode, + } + ] + }, + ) + baseline: Final = ComplexityRouter("test-router", mock_router_instance) + assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2] + assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2] + second_hit_past_the_bound: Final = "yyy " + "a" * 2044 + " zzz" + contribution: Final = ( + router.classify(second_hit_past_the_bound)[1] - baseline.classify(second_hit_past_the_bound)[1] + ) + assert contribution == pytest.approx(0.7 if scoring_mode == "binary" else 0.35) + + @pytest.mark.parametrize( + "prompt,expected_score", + [ + pytest.param("Hello!", 0.0, id="no-hit"), + pytest.param("orbitmesh orbitmesh ORBITMESH again", 0.5, id="one-keyword-repeated"), + pytest.param("create table a; CREATE TABLE b; create table c", 0.5, id="one-pattern-repeated"), + pytest.param("orbitmesh and fluxgate", 1.0, id="two-keywords"), + pytest.param("orbitmesh then create table t", 1.0, id="keyword-plus-pattern"), + pytest.param("create table a; alter table b", 1.0, id="two-patterns"), + pytest.param("orbitmesh fluxgate create table a alter table b", 1.0, id="all-matchers"), + ], + ) + def test_match_count_grades_distinct_matchers( + self, mock_router_instance: MagicMock, prompt: str, expected_score: float + ) -> None: + dimension: Final = { + "name": "graded", + "weight": 0.6, + "keywords": ["orbitmesh", "ORBITMESH", "fluxgate"], + "patterns": [r"\bcreate\s{1,4}table\b", r"\bcreate\s{1,4}table\b", r"\balter\s{1,4}table\b"], + } + baseline: Final = ComplexityRouter("test-router", mock_router_instance) + binary: Final = ComplexityRouter("test-router", mock_router_instance, {"custom_dimensions": [dimension]}) + graded: Final = ComplexityRouter( + "test-router", + mock_router_instance, + {"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]}, + ) + _, baseline_score, baseline_signals = baseline.classify(prompt) + _, binary_score, binary_signals = binary.classify(prompt) + _, graded_score, graded_signals = graded.classify(prompt) + assert graded_score == pytest.approx(baseline_score + 0.6 * expected_score) + assert binary_score == pytest.approx(baseline_score + (0.6 if expected_score else 0.0)) + expected_signals: Final = [*baseline_signals, *(["custom (graded)"] if expected_score else [])] + assert graded_signals == expected_signals + assert binary_signals == expected_signals + + def test_scoring_mode_round_trips_and_defaults_to_binary(self) -> None: + dimension: Final = {"name": "graded", "weight": 0.6, "keywords": ["orbitmesh"]} + legacy: Final = ComplexityRouterConfig.model_validate({"custom_dimensions": [dimension]}) + graded: Final = ComplexityRouterConfig.model_validate( + {"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]} + ) + assert legacy.custom_dimensions[0].scoring_mode == "binary" + assert graded.model_dump(mode="json")["custom_dimensions"][0]["scoring_mode"] == "match_count" + assert ComplexityRouterConfig.model_validate(graded.model_dump(mode="json")) == graded + + def test_custom_dimensions_router_wide_regex_work_is_capped(self) -> None: + heavy: Final = {"weight": 0.5, "patterns": ["a?" * 8 + "z"]} + ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(6)]}) + with pytest.raises(ValidationError, match="regex work estimate is 8939"): + ComplexityRouterConfig.model_validate({"custom_dimensions": [{"name": f"d{i}", **heavy} for i in range(7)]}) + + @pytest.mark.parametrize( + "pattern,work", + [ + pytest.param(r"\b(create|alter|drop)\s{1,4}table\b", 135, id="sql-ddl"), + pytest.param("a?" * 8 + "z", 1277, id="optional-chain-near-cap"), + pytest.param(r"a{0,15}a{0,15}z", 801, id="adjacent-bounded-near-cap"), + pytest.param(r"[a-z0-9_]{3,63}\.(com|net|io)", 1291, id="class-repeat-plus-alternation"), + pytest.param("(?:a|aa)" * 8 + "z", 1787, id="ambiguous-alternation-near-cap"), + pytest.param("a{64}" * 10 + "z", 662, id="long-deterministic-tail"), + ], + ) + def test_custom_pattern_work_stays_cheap_on_adversarial_text( + self, mock_router_instance: MagicMock, pattern: str, work: int + ) -> None: + assert custom_pattern_work(pattern) == work + router: Final = ComplexityRouter( + "test-router", + mock_router_instance, + { + "custom_dimensions": [ + {"name": "bounded", "weight": 0.7, "patterns": [pattern]}, + {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}, + ] + }, + ) + adversarial: Final = "orbitmesh " + "a" * 4000 + started: Final = time.perf_counter() + tier, score, signals = router.classify(adversarial) + elapsed: Final = time.perf_counter() - started + assert signals == ["long (1002 tokens)", "custom (internalFrameworks)"] + assert score == pytest.approx(0.8) + assert tier == ComplexityTier.REASONING + assert elapsed < 0.1 + + class TestAsyncPreRoutingHookEdgeCases: """Test edge cases for async_pre_routing_hook method.""" @@ -1310,6 +1601,7 @@ class TestRouterComplexityDeploymentMethods: def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None: """Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no prompt at all, leaves a router unmetered, so several of them register under a ceiling of one.""" + def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]: llm_config: dict[str, object] = {"model": "gpt-4o-mini"} if preset is not None: @@ -1446,6 +1738,7 @@ class TestRouterComplexityDeploymentMethods: def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None: """tier_labels renames the built-in ladder without defining one, so it stays ungated: two such routers register under a ceiling of one.""" + def labeled(model_name: str, model_id: str) -> dict[str, object]: row = self._router_row(model_name, model_id, "heuristic") row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"} @@ -2319,9 +2612,7 @@ class TestLLMClassifier: assert outcome.classifier_cost == pytest.approx(1.35e-05) @pytest.mark.asyncio - async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks( - self, llm_classifier_config - ): + async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(self, llm_classifier_config): real_router = Router( model_list=[ { @@ -2364,9 +2655,7 @@ class TestLLMClassifier: assert real_router.total_calls["openai/mock-backup-classifier"] == 0 @pytest.mark.asyncio - async def test_aclassify_enforces_total_classifier_deadline( - self, mock_router_instance, llm_classifier_config - ): + async def test_aclassify_enforces_total_classifier_deadline(self, mock_router_instance, llm_classifier_config): cancelled = asyncio.Event() async def slow_classifier(**_kwargs: object) -> None: @@ -3132,11 +3421,11 @@ class TestRouterPreRoutingAliasOverrides: def test_drop_client_effort_carriers_helper_edge_shapes(self): no_pin: Dict = {"thinking": {"type": "adaptive"}} - Router._drop_client_effort_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1}) + Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1}) assert no_pin == {"thinking": {"type": "adaptive"}} non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3} - Router._drop_client_effort_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"}) + Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"}) assert non_dict_carriers == {"output_config": "max", "reasoning": 3} effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}} @@ -3590,11 +3879,11 @@ class TestRouterPreRoutingSharedAliasName: } @staticmethod - async def _routed_call_kwargs(router: Router, **request_params) -> dict: + async def _routed_call_kwargs(router: Router, prompt: str = "hi", **request_params) -> dict: mock_acompletion = AsyncMock(return_value=litellm.ModelResponse(choices=[{"message": {"content": "hi"}}])) with patch.object(litellm, "acompletion", mock_acompletion): await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}], **request_params + model="smart-router", messages=[{"role": "user", "content": prompt}], **request_params ) return mock_acompletion.call_args.kwargs @@ -12213,9 +12502,7 @@ class TestTierHealthFailover: llm_provider="", ) filtered = (*cooling, *blocked, *excluded) - healthy = [ - {"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered - ] + healthy = [{"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered] if not healthy: raise RouterRateLimitError( model=model, cooldown_time=60.0, enable_pre_call_checks=False, cooldown_list=[] @@ -12644,9 +12931,7 @@ class TestTierHealthFailover: assert all(probed is not request_kwargs for probed in router.litellm_router_instance.probed_kwargs) @pytest.mark.asyncio - async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target( - self, mock_router_instance - ): + async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(self, mock_router_instance): """RPM exhaustion is its own verdict from the owner (RouterRateLimitErrorBasic). A peer in that state would be rejected downstream, so it cannot be the substitute.""" from litellm.types.router import RouterRateLimitErrorBasic @@ -12679,9 +12964,7 @@ class TestTierHealthFailover: assert {r.model for r in results} == {"live-c"} @pytest.mark.asyncio - async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces( - self, mock_router_instance - ): + async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces(self, mock_router_instance): """The Responses API carries its prompt as `input`, never as messages. The owner only runs its context-window pre-call check when one of them is present, so dropping `input` would silently skip window filtering on that whole surface.""" @@ -12707,9 +12990,7 @@ class TestTierHealthFailover: ), "the eligibility probe must forward `input` to the owner" @pytest.mark.asyncio - async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target( - self, mock_router_instance - ): + async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(self, mock_router_instance): """The owner answers an unconfigured group with BadRequestError. Reading that as live would both skip failover off it and let it be chosen as a substitute.""" router = self._router( @@ -12898,9 +13179,7 @@ class TestClassifierVision: routed as default_fallback on text the request never contained. """ router = self._router(mock_router_instance, vision={"enabled": True}) - response = await router.async_pre_routing_hook( - model="m", request_kwargs={}, messages=self._turn(IMG_PART) - ) + response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART)) assert response.routing_decision["cause"] == "llm_classifier" assert response.model == "t-complex" assert [block["type"] for block in self._classifier_user_content(mock_router_instance)] == [ @@ -12911,9 +13190,7 @@ class TestClassifierVision: @pytest.mark.asyncio async def test_image_only_turn_still_falls_back_when_vision_is_off(self, mock_router_instance): router = self._router(mock_router_instance, vision={"enabled": False}) - response = await router.async_pre_routing_hook( - model="m", request_kwargs={}, messages=self._turn(IMG_PART) - ) + response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART)) assert response.routing_decision["cause"] == "default_fallback" mock_router_instance.acompletion.assert_not_awaited() @@ -12983,9 +13260,7 @@ class TestClassifierVision: makes the image the only variable; a margin loose enough to leave the score undecided would pass whether or not the guard exists. """ - router = self._router( - mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra - ) + router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra) response = await router.async_pre_routing_hook( model="m", request_kwargs={}, messages=self._turn({"type": "text", "text": "what is this"}, IMG_PART) ) @@ -12999,9 +13274,7 @@ class TestClassifierVision: self, mock_router_instance, classifier_type, extra, short_circuit_cause ): """The negative class: same router, same text, no image, and the scorer still decides.""" - router = self._router( - mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra - ) + router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra) response = await router.async_pre_routing_hook( model="m", request_kwargs={}, messages=[{"role": "user", "content": "what is this"}] ) @@ -13011,3 +13284,379 @@ class TestClassifierVision: def test_max_images_must_be_positive(self): with pytest.raises(ValidationError): ClassifierLLMConfig(model="clf", vision={"enabled": True, "max_images": 0}) + + +class TestMaxTokensFromTierModel: + """The auto-router replaces the caller's output ceiling with the tier model's own, so one + client-side value no longer starves a bigger tier or gets rejected by a smaller one.""" + + COMPLEX_PROMPT: Final = ( + "Design a distributed rate limiter with Redis, sharding and failover. Analyze the consistency " + "tradeoffs and implement the algorithm step by step with tests." + ) + SMALL: Final = { + "model_name": "small", + "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"}, + "model_info": {"max_output_tokens": 8192}, + } + + @staticmethod + def _router( + tier_litellm_params: dict | None = None, + max_tokens_from_tier_model: bool | None = None, + simple_deployments: list[dict] | None = None, + extra_config: dict | None = None, + ) -> Router: + simple_tier: dict = {"model_name": "small"} + if tier_litellm_params: + simple_tier["litellm_params"] = tier_litellm_params + config: dict = { + "tiers": {"SIMPLE": simple_tier, "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"}, + **(extra_config or {}), + } + if max_tokens_from_tier_model is not None: + config["max_tokens_from_tier_model"] = max_tokens_from_tier_model + return Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": {"model": "auto_router/complexity_router", "complexity_router_config": config}, + }, + *(simple_deployments or [TestMaxTokensFromTierModel.SMALL]), + { + "model_name": "big", + "litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "k"}, + "model_info": {"max_output_tokens": 64000}, + }, + ] + ) + + @staticmethod + async def _routed(router: Router, prompt: str = "hi", **request_kwargs) -> dict: + """Drive the real routing entry point and return the request kwargs it leaves behind.""" + deployment = await router.async_get_available_deployment( + model="smart-router", request_kwargs=request_kwargs, messages=[{"role": "user", "content": prompt}] + ) + return {"model": deployment["litellm_params"]["model"], **request_kwargs} + + @staticmethod + async def _routed_responses(router: Router, prompt: str = "hi", **request_kwargs) -> dict: + """The Responses surface hands the router `input` both as the prompt argument and inside the + request kwargs, so the hook sees the same shape the real call carries.""" + routed: dict = {"input": prompt, **request_kwargs} + deployment = await router.async_get_available_deployment( + model="smart-router", request_kwargs=routed, input=prompt + ) + return {"model": deployment["litellm_params"]["model"], **routed} + + @pytest.mark.asyncio + async def test_client_ceiling_is_replaced_by_the_routed_tier_models_ceiling(self): + router = self._router() + + simple = await self._routed(router, max_tokens=8192) + complex_ = await self._routed(router, self.COMPLEX_PROMPT, max_tokens=8192) + + assert (simple["model"], simple["max_tokens"]) == ("anthropic/claude-haiku-4-5", 8192) + assert (complex_["model"], complex_["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + assert "max_output_tokens" not in complex_ + + @pytest.mark.asyncio + async def test_every_client_carrier_of_the_ceiling_is_replaced(self): + sent = await self._routed(self._router(), self.COMPLEX_PROMPT, max_completion_tokens=8192) + + assert sent["max_tokens"] == 64000 + assert "max_completion_tokens" not in sent + + @pytest.mark.asyncio + async def test_responses_surface_gets_the_ceiling_under_its_own_name(self): + sent = await self._routed_responses(self._router(), self.COMPLEX_PROMPT, max_output_tokens=8192) + + assert (sent["model"], sent["max_output_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + assert "max_tokens" not in sent + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "tier_params, responses_call", + [ + ({"max_tokens": 4321}, False), + ({"max_tokens": 4321}, True), + ({"max_completion_tokens": 4321}, False), + ({"max_completion_tokens": 4321}, True), + ({"max_output_tokens": 4321}, False), + ], + ) + async def test_operators_own_tier_ceiling_wins_under_the_surface_name(self, tier_params, responses_call): + router = self._router(tier_litellm_params=tier_params) + if responses_call: + sent = await self._routed_responses(router, max_output_tokens=8192) + else: + sent = await self._routed(router, max_tokens=8192) + + surface_key = "max_output_tokens" if responses_call else "max_tokens" + assert sent[surface_key] == 4321 + assert not (OUTPUT_TOKEN_CEILING_PARAMS - {surface_key}) & sent.keys() + + @pytest.mark.asyncio + async def test_opting_out_forwards_the_client_value_unchanged(self): + sent = await self._routed(self._router(max_tokens_from_tier_model=False), self.COMPLEX_PROMPT, max_tokens=8192) + + assert sent["max_tokens"] == 8192 + + @pytest.mark.asyncio + async def test_a_tier_model_with_an_unknown_ceiling_keeps_the_client_value(self): + unmapped: dict = {"model_name": "small", "litellm_params": {"model": "openai/not-in-any-map", "api_key": "k"}} + + sent = await self._routed(self._router(simple_deployments=[self.SMALL, unmapped]), max_tokens=4000) + + assert sent["max_tokens"] == 4000 + + @pytest.mark.asyncio + async def test_a_multi_deployment_tier_model_uses_its_smallest_ceiling(self): + smaller: dict = { + **self.SMALL, + "litellm_params": {**self.SMALL["litellm_params"], "api_key": "k2"}, + "model_info": {"max_output_tokens": 4096}, + } + + sent = await self._routed(self._router(simple_deployments=[self.SMALL, smaller]), max_tokens=100000) + + assert sent["max_tokens"] == 4096 + + @pytest.mark.asyncio + async def test_ceiling_falls_back_to_the_cost_map(self, monkeypatch): + monkeypatch.setitem( + litellm.model_cost, + "auto-cap-probe-model", + {"litellm_provider": "openai", "mode": "chat", "max_output_tokens": 4242, "max_input_tokens": 100000}, + ) + mapped_only: dict = { + "model_name": "small", + "litellm_params": {"model": "openai/auto-cap-probe-model", "api_key": "k"}, + } + + sent = await self._routed(self._router(simple_deployments=[mapped_only]), max_tokens=8192) + + assert sent["max_tokens"] == 4242 + + @pytest.mark.asyncio + @pytest.mark.parametrize("client_kwargs", [{}, {"max_tokens": 0}], ids=["omitted", "zero"]) + async def test_omitted_and_zero_are_replaced_like_any_other_value(self, client_kwargs): + sent = await self._routed(self._router(), self.COMPLEX_PROMPT, **client_kwargs) + + assert sent["max_tokens"] == 64000 + + @pytest.mark.parametrize( + "tier_params, responses_call, expected", + [ + ({"max_tokens": 1, "temperature": 0.2}, False, {"max_tokens": 1, "temperature": 0.2}), + ({"max_tokens": 1}, True, {"max_output_tokens": 1}), + ({"max_completion_tokens": 2}, False, {"max_tokens": 2}), + ({"max_completion_tokens": 2}, True, {"max_output_tokens": 2}), + ({"max_output_tokens": 3}, False, {"max_tokens": 3}), + ({"max_output_tokens": 3}, True, {"max_output_tokens": 3}), + ({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 1}), + ({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, True, {"max_output_tokens": 3}), + ({"max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 2}), + ({"reasoning_effort": "low"}, True, {"reasoning_effort": "low"}), + ], + ) + def test_every_tier_alias_collapses_onto_the_surface_key(self, tier_params, responses_call, expected): + assert dict(Router._tier_ceiling_under_the_surface_name(tier_params, responses_call=responses_call)) == expected + + @pytest.mark.asyncio + async def test_the_default_fallback_exit_carries_the_ceiling(self): + routed: dict = {"max_tokens": 8192} + deployment = await self._router().async_get_available_deployment( + model="smart-router", request_kwargs=routed, messages=[{"role": "system", "content": "be nice"}] + ) + + assert routed["metadata"]["routing_decision"]["cause"] == "default_fallback" + assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + + @pytest.mark.asyncio + async def test_the_plan_mode_exit_carries_the_ceiling(self): + routed: dict = {"max_tokens": 8192} + deployment = await self._router( + extra_config={"plan_mode_min_tier": "REASONING"} + ).async_get_available_deployment( + model="smart-router", + request_kwargs=routed, + messages=[ + {"role": "user", "content": "plan the refactor"}, + {"role": "system", "content": "Plan mode is active"}, + ], + ) + + assert routed["metadata"]["routing_decision"]["cause"] == "plan_mode" + assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + + @pytest.mark.asyncio + async def test_a_default_model_landing_with_no_tier_still_gets_its_ceiling(self): + strategy = ComplexityRouter( + model_name="smart-router", + litellm_router_instance=self._router(), + complexity_router_config={"tiers": {"SIMPLE": "small"}, "default_model": "big"}, + ) + + assert dict(strategy._litellm_params_for_model(None, "big")) == {"max_tokens": 64000} + + @pytest.mark.asyncio + async def test_a_fallback_into_a_plain_group_gets_the_callers_ceiling_back(self): + """A model-group fallback re-enters routing with the same kwargs; a Sonnet-sized ceiling + must not ride onto the plain group the caller configured as the fallback.""" + big: dict = { + "model_name": "big", + "litellm_params": { + "model": "anthropic/claude-sonnet-5", + "api_key": "k", + "mock_response": "litellm.InternalServerError", + }, + "model_info": {"max_output_tokens": 64000}, + } + plain: dict = { + "model_name": "plain", + "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k", "mock_response": "ok"}, + } + router = Router( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "big", "MEDIUM": "big", "COMPLEX": "big", "REASONING": "big"} + }, + }, + }, + big, + plain, + ], + fallbacks=[{"smart-router": ["plain"]}], + num_retries=0, + ) + recorder = _OutputCeilingRecorder() + litellm.callbacks.append(recorder) + try: + await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": self.COMPLEX_PROMPT}], max_tokens=8192 + ) + finally: + litellm.callbacks.remove(recorder) + + assert recorder.seen == [("claude-sonnet-5", 64000), ("claude-haiku-4-5", 8192)] + + @pytest.mark.asyncio + async def test_a_caller_seeded_stamp_cannot_inject_kwargs_on_a_plain_group(self): + """The stamp sits in a metadata bucket a caller can write; a planted one must yield + nothing but integer ceiling carriers, never a redirected api_base or credential.""" + planted: dict = { + "api_base": "https://attacker.example", + "api_key": "stolen", + "max_tokens": "not-an-int", + "max_completion_tokens": True, + "max_output_tokens": 321, + } + routed: dict = {"max_tokens": 8192, "metadata": {"_client_output_ceiling": planted}} + + await self._router().async_get_available_deployment( + model="big", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}] + ) + + assert {k: v for k, v in routed.items() if k not in ("metadata", "model_info")} == {"max_output_tokens": 321} + + @pytest.mark.asyncio + async def test_the_pass_through_routing_entry_point_pins_and_restores_the_same_way(self): + pass_through: dict = {**self.SMALL["litellm_params"], "use_in_pass_through": True} + small: dict = {**self.SMALL, "litellm_params": pass_through} + plain: dict = {**small, "model_name": "plain"} + router = self._router(simple_deployments=[small, plain]) + for deployment in router.model_list: + deployment["litellm_params"]["use_in_pass_through"] = True + routed: dict = {"max_tokens": 8192} + + deployment = await router.async_get_available_deployment_for_pass_through( + model="smart-router", request_kwargs=routed, messages=[{"role": "user", "content": self.COMPLEX_PROMPT}] + ) + pinned = routed["max_tokens"] + await router.async_get_available_deployment_for_pass_through( + model="plain", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}] + ) + + assert (deployment["litellm_params"]["model"], pinned, routed["max_tokens"]) == ( + "anthropic/claude-sonnet-5", + 64000, + 8192, + ) + + @pytest.mark.asyncio + async def test_the_classifier_fallback_exit_carries_the_ceiling(self): + router = self._router( + extra_config={ + "classifier_type": "llm", + "classifier_llm_config": {"model": "no-such-classifier", "timeout_ms": 400}, + "classifier_fallback": "default_model", + "default_model": "big", + } + ) + routed: dict = {"max_tokens": 8192} + + deployment = await router.async_get_available_deployment( + model="smart-router", request_kwargs=routed, messages=[{"role": "user", "content": "hi"}] + ) + + assert routed["metadata"]["routing_decision"]["cause"] == "default_model_fallback" + assert (deployment["litellm_params"]["model"], routed["max_tokens"]) == ("anthropic/claude-sonnet-5", 64000) + + @pytest.mark.parametrize( + "value, expected", + [(8192, 8192), ("8192", 8192), (100.9, 100), (0, 0), (-1, None), (True, None), ("x", None), (None, None)], + ) + def test_a_client_cap_is_read_as_an_integer_or_ignored(self, value, expected): + assert as_output_cap(value) == expected + + def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self): + stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}} + Router._restore_client_ceiling_no_tier_pins(stamped) + assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192} + + coerced: dict = { + "max_tokens": 64000, + "metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}}, + } + Router._restore_client_ceiling_no_tier_pins(coerced) + assert {k: v for k, v in coerced.items() if k != "metadata"} == { + "max_tokens": 8192, + "max_completion_tokens": 100, + } + + unstamped: dict = {"max_tokens": 64000, "metadata": {}} + Router._restore_client_ceiling_no_tier_pins(unstamped) + assert unstamped["max_tokens"] == 64000 + + @pytest.mark.asyncio + async def test_pinning_stamps_the_callers_carriers_once(self): + router = self._router() + request_kwargs: dict = {"max_completion_tokens": 8192} + + first = router._pin_tier_params_onto_request( + model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False + ) + second = router._pin_tier_params_onto_request( + model="big", tier_litellm_params={"max_tokens": 32000}, request_kwargs=request_kwargs, responses_call=False + ) + none = router._pin_tier_params_onto_request( + model="big", tier_litellm_params=None, request_kwargs=request_kwargs, responses_call=False + ) + + assert (first, second, none) == (True, True, False) + assert request_kwargs["max_tokens"] == 32000 + assert request_kwargs["metadata"]["_client_output_ceiling"] == {"max_completion_tokens": 8192} + + +class _OutputCeilingRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen: list[tuple[str, int | None]] = [] + + def log_pre_api_call(self, model, messages, kwargs): + self.seen.append((model, kwargs.get("optional_params", {}).get("max_tokens"))) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py new file mode 100644 index 00000000000..9efa526fc02 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -0,0 +1,187 @@ +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler + +GROUP: Final = "least-busy-group" +DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} +DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} +HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] + + +def _call_kwargs(deployment_id: str) -> dict[str, object]: + return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} + + +class SharedRedisCounters: + """Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when + the key is created, and ordered reads that raise rather than invent a value.""" + + def __init__(self) -> None: + self.counts: dict[str, int] = {} + self.ttls: dict[str, int] = {} + + def count(self, key: str) -> int | None: + return self.counts.get(key) + + def expire(self, key: str) -> None: + self.counts.pop(key, None) + self.ttls.pop(key, None) + + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + incremented: Final = max(0, self.counts.get(key, 0) + value) + self.counts[key] = incremented + self.ttls.setdefault(key, ttl) + return incremented + + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + return self.increment_with_floor(key, value, ttl) + + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return tuple(self.counts.get(key) for key in key_list) + + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return self.batch_get_counts(key_list) + + +def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: + cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double + return LeastBusyLoggingHandler(router_cache=cache) + + +@pytest.mark.asyncio +async def test_worker_routes_around_a_request_another_worker_started() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + await picking_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await streaming_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_sync_pick_reads_the_shared_counts() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + picking_worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + streaming_worker.log_failure_event(_call_kwargs("dep-a"), None, None, None) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: + """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that + stuck count an hour after the key was created, which only works while nothing writes the TTL + again: a handler that refreshed it on every touch would keep the count alive for as long as + the group takes traffic, and the deployment would read busier than it is forever.""" + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + key: Final = f"{GROUP}_request_count:dep-a" + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + + shared.ttls[key] = 5 + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(key) == 1 + assert shared.ttls == {key: 5} + + +@pytest.mark.asyncio +async def test_counts_stay_in_memory_without_redis() -> None: + worker: Final = _worker(None) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + +class UnavailableRedis(SharedRedisCounters): + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + raise ConnectionError("redis is down") + + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + raise ConnectionError("redis is down") + + +@pytest.mark.asyncio +async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> None: + worker: Final = _worker(UnavailableRedis()) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + shared.expire(f"{GROUP}_request_count:dep-a") + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 1 + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +@pytest.mark.asyncio +async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() -> None: + worker: Final = _worker(None) + in_memory: Final = worker.router_cache.in_memory_cache + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + in_memory.delete_cache(f"{GROUP}_request_count:dep-a") + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +def test_calls_without_a_deployment_are_ignored() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) + worker.log_pre_api_call(model="m", messages=[], kwargs={}) + + assert shared.counts == {} diff --git a/tests/test_litellm/router_strategy/test_lowest_latency.py b/tests/test_litellm/router_strategy/test_lowest_latency.py index eb02459be68..812d7bbff32 100644 --- a/tests/test_litellm/router_strategy/test_lowest_latency.py +++ b/tests/test_litellm/router_strategy/test_lowest_latency.py @@ -165,6 +165,133 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): json.dumps({"latency": latencies}) +MODEL_GROUP = "gpt-4o-mini" +FAST_TTFT_ID = "fast-ttft-short-output" +SLOW_TTFT_ID = "slow-ttft-long-output" +STREAMING_DEPLOYMENTS = [ + {"model_info": {"id": FAST_TTFT_ID}, "litellm_params": {}}, + {"model_info": {"id": SLOW_TTFT_ID}, "litellm_params": {}}, +] + + +def _streaming_kwargs(deployment_id: str, start_time: datetime, ttft_seconds: float): + return { + "litellm_params": { + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": deployment_id}, + }, + "stream": True, + "completion_start_time": start_time + timedelta(seconds=ttft_seconds), + } + + +def _recorded_ttft(cache: DualCache, deployment_id: str): + cached = cache.get_cache(key=f"{MODEL_GROUP}_map") or {} + return cached.get(deployment_id, {}).get("time_to_first_token_seconds", []) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_streaming_ttft_ranking_ignores_completion_length(sync_mode: bool): + """Deployment A: TTFT 1s, 50 completion tokens. Deployment B: TTFT 3s, 500 + completion tokens. Dividing TTFT by completion tokens made B look faster + (3/500 = 0.006 beats 1/50 = 0.02); actual TTFT must win.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + start_time = datetime(2026, 1, 1, 12, 0, 0) + end_time = start_time + timedelta(seconds=10) + + samples = ( + (FAST_TTFT_ID, 1.0, 50), + (SLOW_TTFT_ID, 3.0, 500), + ) + for deployment_id, ttft, completion_tokens in samples: + kwargs = _streaming_kwargs(deployment_id, start_time, ttft) + response_obj = _chat_response(completion_tokens=completion_tokens) + if sync_mode: + handler.log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time + ) + else: + await handler.async_log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=end_time + ) + + assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(1.0)] + assert _recorded_ttft(cache, SLOW_TTFT_ID) == [pytest.approx(3.0)] + + request_kwargs = {"stream": True, "metadata": {}} + if sync_mode: + picked = handler.get_available_deployments( + model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs + ) + else: + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, healthy_deployments=STREAMING_DEPLOYMENTS, request_kwargs=request_kwargs + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_ttft_window_keeps_newest_samples_when_full(sync_mode: bool): + """Float timestamps, as the SDK passes them. Once max_latency_list_size + samples exist the oldest TTFT is dropped so the window slides.""" + max_size = 3 + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": max_size}) + start_time = 1_700_000_000.0 + ttfts = (0.1, 0.2, 0.3, 0.4) + + for ttft in ttfts: + kwargs = { + "litellm_params": { + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": FAST_TTFT_ID}, + }, + "stream": True, + "completion_start_time": start_time + ttft, + } + response_obj = _chat_response(completion_tokens=1) + if sync_mode: + handler.log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0 + ) + else: + await handler.async_log_success_event( + response_obj=response_obj, kwargs=kwargs, start_time=start_time, end_time=start_time + 1.0 + ) + + assert _recorded_ttft(cache, FAST_TTFT_ID) == [pytest.approx(ttft) for ttft in ttfts[-max_size:]] + + +@pytest.mark.asyncio +async def test_streaming_routing_ignores_per_token_ttft_samples_from_older_workers(): + """Workers on the previous release share the Redis map and keep writing + seconds-per-token under the old "time_to_first_token" key during a rolling + deploy. Those samples favor SLOW; routing must only read the seconds key.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + cache.set_cache( + key=f"{MODEL_GROUP}_map", + value={ + FAST_TTFT_ID: {"time_to_first_token": [0.02], "time_to_first_token_seconds": [1.0]}, + SLOW_TTFT_ID: {"time_to_first_token": [0.006], "time_to_first_token_seconds": [3.0]}, + }, + ) + + picked = await handler.async_get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=STREAMING_DEPLOYMENTS, + request_kwargs={"stream": True, "metadata": {}}, + ) + + assert picked is not None + assert picked["model_info"]["id"] == FAST_TTFT_ID + + @pytest.mark.asyncio @pytest.mark.parametrize( "cached_entry", diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 5599c5aad63..5f37842305d 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import RoutingGroup, RoutingStrategy @@ -435,6 +436,81 @@ def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeyp assert router._group_selectors == {} +def test_two_least_busy_groups_count_a_request_once(monkeypatch): + """ + Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and + back down from the success hooks on `litellm.callbacks`. The success list drops a second + selector of the same class, so a pre-call list that kept both counted every request twice + and released it once, and the deployment's in-flight count climbed until it looked pinned. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + router = _build_router( + routing_strategy="least-busy", + routing_groups=[ + { + "group_name": "fast", + "models": ["filtered-model"], + "routing_strategy": "least-busy", + } + ], + ) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + +def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): + """ + Least-busy hangs its counting off litellm's global callback lists, and those lists keep one + logger per class unless the instances differ in a plain attribute. Two routers in one process + (a second Router, or a per-request `user_config` one) therefore have to register separately: + a second router whose selector is dropped counts nothing, reads zero for every deployment, + and sends every request to whichever one is listed first. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + first = _build_router(routing_strategy="least-busy") + second = _build_router(routing_strategy="least-busy") + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 1 + assert ( + second.get_available_deployment(model="filtered-model", messages=[])["model_info"]["id"] + == "deploy-2" + ) + + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert first.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + # --------------------------------------------------------------------------- # Direct helper coverage # --------------------------------------------------------------------------- @@ -716,15 +792,168 @@ def test_strategy_reinit_unregisters_override_selectors(): router = _build_router(routing_strategy="least-busy") override_selector = router._get_override_strategy_selector("latency-based-routing") assert override_selector is not None - assert any(id(cb) == id(override_selector) for cb in litellm.callbacks) + assert not any(cb is override_selector for cb in litellm.callbacks) router.update_settings(routing_strategy="latency-based-routing") assert router._override_selectors == {} - assert not any(id(cb) == id(override_selector) for cb in litellm.callbacks) + assert not any(cb is override_selector for cb in litellm.callbacks) assert router._get_override_strategy_selector("latency-based-routing") is router.lowestlatency_logger +def test_override_selectors_are_not_registered_process_wide(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + router = _build_router(routing_strategy="simple-shuffle") + v2_selector = router._get_override_strategy_selector("usage-based-routing-v2") + least_busy_selector = router._get_override_strategy_selector("least-busy") + assert v2_selector is not None and least_busy_selector is not None + + assert litellm.callbacks == [] + assert litellm.input_callback == [] + + +def _rpm_limited_model_list(): + return [ + { + "model_name": "other-model", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-test-3", + "api_base": "https://example.invalid", + "rpm": 1, + }, + "model_info": {"id": "deploy-3"}, + }, + ] + + +async def _mock_completion(router, **override): + return await router.acompletion( + model="other-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok", **override + ) + + +@pytest.mark.asyncio +async def test_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(): + router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) + + await _mock_completion(router, routing_strategy="usage-based-routing-v2") + override_selector = router._override_selectors["usage-based-routing-v2"] + assert not any(cb is override_selector for cb in litellm.callbacks) + + with patch.object( + override_selector, "async_pre_call_check", wraps=override_selector.async_pre_call_check + ) as pre_call_spy: + for _ in range(2): + plain = await _mock_completion(router) + assert plain.choices[0].message.content == "ok" + assert not pre_call_spy.called + + with pytest.raises(litellm.RateLimitError): + await _mock_completion(router, routing_strategy="usage-based-routing-v2") + + +def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(): + router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) + messages = [{"role": "user", "content": "hi"}] + + router.completion( + model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2" + ) + override_selector = router._override_selectors["usage-based-routing-v2"] + assert not any(cb is override_selector for cb in litellm.callbacks) + + for _ in range(2): + plain = router.completion(model="other-model", messages=messages, mock_response="ok") + assert plain.choices[0].message.content == "ok" + + with pytest.raises(ValueError, match="No deployments available"): + router.completion( + model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2" + ) + + +@pytest.mark.asyncio +async def test_override_selector_pre_call_check_only_runs_for_override_selectors(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", []) + deployment = _rpm_limited_model_list()[0] + + override_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle") + override_selector = override_router._get_override_strategy_selector("usage-based-routing-v2") + await override_router._async_override_selector_pre_call_check( + "usage-based-routing-v2", override_selector, deployment, None + ) + with pytest.raises(litellm.RateLimitError): + override_router._override_selector_pre_call_check("usage-based-routing-v2", override_selector, deployment) + + default_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="usage-based-routing-v2") + for _ in range(2): + await default_router._async_override_selector_pre_call_check( + "usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment, None + ) + default_router._override_selector_pre_call_check( + "usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment + ) + await default_router._async_override_selector_pre_call_check(None, None, deployment, None) + default_router._override_selector_pre_call_check(None, None, deployment) + + +@pytest.mark.asyncio +async def test_usage_based_v2_override_enforces_rpm_when_a_specific_deployment_is_requested(): + router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) + kwargs = {"model": "deploy-3", "messages": [{"role": "user", "content": "hi"}], "mock_response": "ok"} + + first = await router.acompletion(**kwargs, routing_strategy="usage-based-routing-v2") + assert first.choices[0].message.content == "ok" + with pytest.raises(litellm.RateLimitError): + await router.acompletion(**kwargs, routing_strategy="usage-based-routing-v2") + assert (await router.acompletion(**kwargs)).choices[0].message.content == "ok" + + +def test_sync_usage_based_v2_override_enforces_rpm_when_a_specific_deployment_is_requested(): + router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) + kwargs = {"model": "deploy-3", "messages": [{"role": "user", "content": "hi"}], "mock_response": "ok"} + + first = router.completion(**kwargs, routing_strategy="usage-based-routing-v2") + assert first.choices[0].message.content == "ok" + with pytest.raises(litellm.RateLimitError): + router.completion(**kwargs, routing_strategy="usage-based-routing-v2") + assert router.completion(**kwargs).choices[0].message.content == "ok" + + +def _pass_through_rpm_limited_model_list(): + deployment = _rpm_limited_model_list()[0] + return [{**deployment, "litellm_params": {**deployment["litellm_params"], "use_in_pass_through": True}}] + + +@pytest.mark.asyncio +async def test_async_early_return_paths_run_the_override_pre_call_check(): + router = Router(model_list=_pass_through_rpm_limited_model_list(), routing_strategy="simple-shuffle") + override = {"routing_strategy": "usage-based-routing-v2"} + + pinned = await router.async_get_available_deployment( + model="other-model", request_kwargs={**override, "_encrypted_content_affinity_pinned": True} + ) + assert pinned["model_info"]["id"] == "deploy-3" + with pytest.raises(litellm.RateLimitError): + await router.async_get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override) + plain = await router.async_get_available_deployment_for_pass_through(model="deploy-3", request_kwargs={}) + assert plain["model_info"]["id"] == "deploy-3" + + +def test_sync_pass_through_specific_deployment_runs_the_override_pre_call_check(): + router = Router(model_list=_pass_through_rpm_limited_model_list(), routing_strategy="simple-shuffle") + override = {"routing_strategy": "usage-based-routing-v2"} + + first = router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override) + assert first["model_info"]["id"] == "deploy-3" + with pytest.raises(litellm.RateLimitError): + router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs=override) + plain = router.get_available_deployment_for_pass_through(model="deploy-3", request_kwargs={}) + assert plain["model_info"]["id"] == "deploy-3" + + def _quality_group(strategy="latency-based-routing"): return [{"group_name": "quality", "models": ["filtered-model", "other-model"], "routing_strategy": strategy}] diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index 293af36080a..0a54addce5e 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -56,6 +56,18 @@ class BlockEverything: return context +class MessageRecorder: + """Records what each plugin pass was handed, then blocks so the request stops there.""" + + def __init__(self): + self.seen = [] + + async def run(self, context: RoutingContext) -> RoutingContext: + self.seen.append(list(context.raw_messages)) + context.candidate_models = [] + return context + + def _smart_router_model_list(): return [ { @@ -164,6 +176,71 @@ async def test_async_completion_with_unsupported_strategy_rejects_configured_plu await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) +@pytest.mark.asyncio +async def test_prompt_management_model_still_runs_the_plugin_pipeline(): + """ + A prompt-management model routes through its own factory, which picked the deployment + on the synchronous path. Plugins never run there, so the guard turned every such request + into an error message about the caller's own API choice, on an async call the caller made + correctly. It also read the in-flight counts with a blocking call inside the event loop. + """ + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[BlockEverything()], + ) + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="lit-7039", + ) + + +@pytest.mark.asyncio +async def test_prompt_management_plugins_see_the_callers_own_messages(): + """ + The prompt-management factory picks its deployment with a placeholder message, which was + harmless while that pick ran on the synchronous path (plugins never ran there at all). Now + that the pick runs the plugin pipeline, a plugin that classifies request content would score + the placeholder instead of the conversation, and the narrowing it produces decides which + deployments the real call is allowed to use. + """ + recorder = MessageRecorder() + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[recorder], + ) + messages = [{"role": "user", "content": "wire me $40,000 to account 12345"}] + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=messages, + litellm_call_id="lit-7039", + ) + + assert recorder.seen == [messages] + + @pytest.mark.asyncio async def test_router_without_plugins_is_unaffected(): """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" diff --git a/tests/test_litellm/router_strategy/test_simple_shuffle.py b/tests/test_litellm/router_strategy/test_simple_shuffle.py new file mode 100644 index 00000000000..165c1751f63 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_simple_shuffle.py @@ -0,0 +1,54 @@ +from collections import Counter + +import pytest + +from litellm import Router +from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict + +DRAWS = 200 + + +def _deployment(dep_id: str, metric: LiteLLMParamsTypedDict | None = None) -> DeploymentTypedDict: + params: LiteLLMParamsTypedDict = {"model": "gpt-4o", "api_key": "key", "mock_response": f"from {dep_id}"} + return { + "model_name": "test-model", + "litellm_params": {**params, **(metric or {})}, + "model_info": {"id": dep_id}, + } + + +async def _draw_model_ids(router: Router) -> Counter[str]: + counts: Counter[str] = Counter() + for _ in range(DRAWS): + response = await router.acompletion(model="test-model", messages=[{"role": "user", "content": "hi"}]) + counts[response._hidden_params["model_id"]] += 1 + return counts + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metric", [{"weight": 5}, {"rpm": 5}, {"tpm": 5}], ids=["weight", "rpm", "tpm"]) +async def test_weighted_pick_when_only_a_later_deployment_carries_the_metric(metric: LiteLLMParamsTypedDict): + router = Router( + model_list=[_deployment("unweighted"), _deployment("weighted", metric)], + routing_strategy="simple-shuffle", + num_retries=0, + ) + + counts = await _draw_model_ids(router) + + assert counts["weighted"] == DRAWS + assert counts["unweighted"] == 0 + + +@pytest.mark.asyncio +async def test_uniform_pick_when_every_configured_weight_is_zero(): + router = Router( + model_list=[_deployment("unweighted"), _deployment("standby", {"weight": 0})], + routing_strategy="simple-shuffle", + num_retries=0, + ) + + counts = await _draw_model_ids(router) + + assert counts["unweighted"] > 0 + assert counts["standby"] > 0 diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 79ae00e155c..030bdfe03e9 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -332,6 +332,67 @@ async def test_per_request_enable_prompt_caching_reaches_the_affinity_key(monkey assert filtered == [deployments[1]] +@pytest.mark.asyncio +async def test_claude_code_one_shot_subagent_does_not_reuse_an_auto_injected_affinity_key(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + messages = cast(List[AllMessageValues], [{"role": "user", "content": "unique " * 3000}]) + request_kwargs = { + "system": [ + { + "type": "text", + "text": "x-anthropic-billing-header: cc_version=2.1.263; cc_is_subagent=true;", + } + ], + "proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.263 (external, cli)"}}, + } + auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=(AUTO_CACHING_MODEL,), + ) + assert auto_injected_messages != messages + await PromptCachingCache(cache=cache).async_add_model_id( + model_id="dep-2", messages=auto_injected_messages, tools=None + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs=request_kwargs, + ) + + assert filtered == deployments + + +@pytest.mark.asyncio +async def test_root_cache_control_does_not_reuse_an_auto_injected_affinity_key(monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + messages = _auto_caching_messages() + auto_injected_messages = AnthropicCacheControlHook.messages_with_default_injections( + messages=messages, + models=(AUTO_CACHING_MODEL,), + ) + assert auto_injected_messages != messages + await PromptCachingCache(cache=cache).async_add_model_id( + model_id="dep-2", messages=auto_injected_messages, tools=None + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs={"cache_control": {"type": "ephemeral"}}, + ) + + assert filtered == deployments + + @pytest.mark.asyncio async def test_tool_marked_cache_control_keeps_routing_off_another_requests_prefix(monkeypatch, local_model_cost_map): """ 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..75c115cd3ab 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 @@ -3,9 +3,11 @@ from __future__ import annotations from collections.abc import Mapping +from typing import Final import pytest +from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig from litellm.router_utils.auto_router_tuning_baseline import ( DEFAULT_TUNING_FINGERPRINT, HEURISTIC_V1_TUNING_FIELDS, @@ -20,6 +22,33 @@ from litellm.router_utils.auto_router_tuning_baseline import ( _TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"} _ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"} +_KEYWORD_DIMENSION: Final = {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]} +_HISTORICAL_FINGERPRINTS: Final = ( + ({}, "44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a"), + ( + {"custom_dimensions": [_KEYWORD_DIMENSION]}, + "b5c3c3f3be6341a8a16148d68d9067e03f94ed01a0bbfcde7955763042744372", + ), + ( + {"custom_dimensions": [{"name": "sqlDdl", "weight": 0.4, "patterns": [r"\bCREATE\s{1,4}TABLE\b"]}]}, + "814ce0017fc7f60a160b262f658d910e9bdf784e6139a4ba4f1e2657aa203950", + ), + ( + { + "tiers": _TIERS, + "dimension_weights": {"codePresence": 0.3}, + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.2, + "keywords": ["orbitmesh", "fluxgate"], + "patterns": [r"\bALTER\s{1,4}TABLE\b"], + } + ], + }, + "38970dc9224e265ab38c89674563d8d0537822591f9239b45251db6f5ca6cc39", + ), +) def _router( @@ -58,6 +87,7 @@ class TestTuningFingerprint: "reasoning_override_min_score": 0.05, "token_thresholds": {"simple": 20, "complex": 500}, "dimension_weights": {"codePresence": 0.9}, + "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}], "code_keywords": ["orionflow"], "reasoning_keywords": ["deduce"], "technical_keywords": ["ledgerkit"], @@ -72,6 +102,30 @@ 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 + + @pytest.mark.parametrize(("config", "fingerprint"), _HISTORICAL_FINGERPRINTS) + def test_fingerprints_recorded_before_scoring_mode_existed_are_preserved( + self, config: Mapping[str, object], fingerprint: str + ) -> None: + """Literal hashes captured from the merged implementation at 9bc9104102, before CustomDimension.scoring_mode.""" + assert tuning_fingerprint(config) == fingerprint + + def test_binary_scoring_mode_hashes_like_its_absence(self) -> None: + historical: Final = tuning_fingerprint({"custom_dimensions": [_KEYWORD_DIMENSION]}) + explicit: Final = tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "binary"}]}) + reserialized: Final = ComplexityRouterConfig.model_validate( + {"custom_dimensions": [_KEYWORD_DIMENSION]} + ).model_dump(mode="json", include={"custom_dimensions"}) + assert reserialized["custom_dimensions"][0]["scoring_mode"] == "binary" + assert reserialized["custom_dimensions"][0]["patterns"] == [] + assert historical == explicit == tuning_fingerprint(reserialized) + assert ( + tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]}) + != historical + ) + def test_tier_model_overrides_change_the_fingerprint(self) -> None: plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}}) with_override = tuning_fingerprint( @@ -80,11 +134,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 +243,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 +262,42 @@ 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 + ) + + @pytest.mark.parametrize( + "edit", + [ + pytest.param({"weight": 0.9}, id="weight"), + pytest.param({"scoring_mode": "match_count"}, id="scoring-mode"), + ], + ) + def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self, edit: Mapping[str, object]) -> None: + baselines: Final = snapshot_tuning_baselines(()) + original: Final = _router("a", {}) + config: Final = {"custom_dimensions": [_KEYWORD_DIMENSION]} + edited_config: Final = {"custom_dimensions": [{**_KEYWORD_DIMENSION, **edit}]} + added: Final = _router("a", config) + edited: Final = _router("a", edited_config) + second: Final = _router("b", config) + + assert tuning_fingerprint(config) != tuning_fingerprint(edited_config) + assert mutable_tuned_identities((added,), baselines) == {router_identity(original)} + assert tuning_quota_violation(candidate=added, others=(original,), baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=edited, others=(added,), baselines=baselines, limit=1) is None + assert tuning_quota_violation(candidate=second, others=(edited,), baselines=baselines, limit=1) is not None + assert tuning_quota_violation(candidate=original, others=(edited,), baselines=baselines, limit=1) is None + assert mutable_tuned_identities((original,), baselines) == frozenset() + assert tuning_quota_violation(candidate=second, others=(original,), baselines=baselines, limit=1) is None + + def test_graded_dimension_recorded_at_snapshot_is_its_own_baseline(self) -> None: + graded: Final = _router("a", {"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]}) + baselines: Final = snapshot_tuning_baselines((graded,)) + assert mutable_tuned_identities((graded,), baselines) == frozenset() + reverted_to_binary: Final = _router("a", {"custom_dimensions": [_KEYWORD_DIMENSION]}) + assert mutable_tuned_identities((reverted_to_binary,), baselines) == {router_identity(graded)} 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_cooldown_cache.py b/tests/test_litellm/router_utils/test_cooldown_cache.py index 68e9aeaa4fc..6f90fa8465f 100644 --- a/tests/test_litellm/router_utils/test_cooldown_cache.py +++ b/tests/test_litellm/router_utils/test_cooldown_cache.py @@ -268,12 +268,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired cooldown entry must not appear in active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" + assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache" def test_active_entry_is_returned(self): """ @@ -289,7 +289,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -312,14 +312,14 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - (60.0 - remaining), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + cc.in_memory_cache.set_cache(key, value, ttl=600) - before_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + before_expiry = cc.in_memory_cache.ttl_dict.get(key) assert before_expiry is not None cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry is not None corrected_remaining = after_expiry - time.time() assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" @@ -340,12 +340,12 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time() - 120.0, "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, expired_value, ttl=600) + cc.in_memory_cache.set_cache(key, expired_value, ttl=600) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) assert active == [], "Expired entry must not appear in async active cooldowns" - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None @pytest.mark.asyncio async def test_async_active_entry_is_returned(self): @@ -363,7 +363,7 @@ class TestCooldownCacheTTLCorrection: "timestamp": time.time(), "cooldown_time": 60.0, } - cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + cc.in_memory_cache.set_cache(key, active_value, ttl=60) active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) @@ -389,18 +389,18 @@ class TestCorrectedActiveCooldown: cc = self._make_cooldown_cache() key = "deployment:expired-dep:cooldown" entry = self._entry(timestamp=time.time() - 120.0, cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=600) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) assert result is None - assert cc.cache.in_memory_cache.get_cache(key) is None + assert cc.in_memory_cache.get_cache(key) is None def test_active_entry_within_window_returns_value(self): cc = self._make_cooldown_cache() key = "deployment:active-dep:cooldown" entry = self._entry(timestamp=time.time(), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=60) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) @@ -412,12 +412,12 @@ class TestCorrectedActiveCooldown: key = "deployment:backfilled-dep:cooldown" remaining = 30.0 entry = self._entry(timestamp=time.time() - (60.0 - remaining), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=600) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=600) result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) assert result is not None - corrected_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + corrected_expiry = cc.in_memory_cache.ttl_dict.get(key) assert corrected_expiry is not None assert corrected_expiry - time.time() <= 60.0 @@ -425,10 +425,160 @@ class TestCorrectedActiveCooldown: cc = self._make_cooldown_cache() key = "deployment:normal-dep:cooldown" entry = self._entry(timestamp=time.time(), cooldown_time=60.0) - cc.cache.in_memory_cache.set_cache(key, dict(entry), ttl=60) - original_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + cc.in_memory_cache.set_cache(key, dict(entry), ttl=60) + original_expiry = cc.in_memory_cache.ttl_dict.get(key) cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) - after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + after_expiry = cc.in_memory_cache.ttl_dict.get(key) assert after_expiry == original_expiry + + +class SharedRedisDouble: + """ + In-process stand-in for RedisCache, shared by several DualCache instances so that + tests can model two proxy replicas talking to one Redis. + """ + + def __init__(self) -> None: + self.store: dict = {} # mutable-ok: stands in for Redis' own mutable keyspace + + def set_cache(self, key, value, **kwargs): + self.store[key] = value + + async def async_set_cache(self, key, value, **kwargs): + self.store[key] = value + + def batch_get_cache(self, key_list, parent_otel_span=None, **kwargs): + return {key: self.store.get(key) for key in key_list} + + async def async_batch_get_cache(self, key_list, parent_otel_span=None, **kwargs): + return {key: self.store.get(key) for key in key_list} + + +class TestCooldownPropagationBetweenReplicas: + """ + A cooldown written by one replica has to reach its siblings quickly. The router's own + DualCache re-reads a key that is missing from memory only every 10s, so cooldown reads + get their own cache with a much shorter Redis read interval. + """ + + def _make_replica(self, redis: SharedRedisDouble, read_interval: float | None = None) -> CooldownCache: + router_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis) + if read_interval is None: + return CooldownCache(cache=router_cache, default_cooldown_time=60.0) + return CooldownCache( + cache=router_cache, + default_cooldown_time=60.0, + redis_read_interval_seconds=read_interval, + ) + + @pytest.mark.asyncio + async def test_sibling_replica_sees_cooldown_within_configured_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis, read_interval=0.25) + replica_b = self._make_replica(redis, read_interval=0.25) + model_id = "shared-deployment" + + assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(0.3) + + active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "sibling replica must pick up a cooldown written by another replica within the read interval" + ) + + @pytest.mark.asyncio + async def test_sibling_replica_sees_cooldown_within_default_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis) + replica_b = self._make_replica(redis) + model_id = "default-interval-deployment" + + assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(1.2) + + active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "the shipped default read interval must let a sibling replica see a cooldown about a second later" + ) + + def test_sync_read_path_sees_sibling_cooldown_within_read_interval(self): + redis = SharedRedisDouble() + replica_a = self._make_replica(redis, read_interval=0.25) + replica_b = self._make_replica(redis, read_interval=0.25) + model_id = "sync-shared-deployment" + + assert replica_b.get_active_cooldowns([model_id], parent_otel_span=None) == [] + + replica_a.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + time.sleep(0.3) + + active = replica_b.get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active] + + @pytest.mark.asyncio + async def test_redis_attached_after_construction_is_still_used(self): + redis = SharedRedisDouble() + router_cache = DualCache(in_memory_cache=InMemoryCache()) + writer = CooldownCache(cache=router_cache, default_cooldown_time=60.0, redis_read_interval_seconds=0.25) + router_cache.attach_redis_cache(redis) + reader = self._make_replica(redis, read_interval=0.25) + model_id = "late-redis-deployment" + + writer.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=60.0, + ) + + active = await reader.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "a router that wires Redis after building its cooldown cache must still publish cooldowns to it" + ) + + +class TestCooldownSurvivesUnrelatedCacheTraffic: + @pytest.mark.asyncio + async def test_unrelated_router_cache_writes_do_not_evict_active_cooldown(self): + router_cache = DualCache(in_memory_cache=InMemoryCache()) + cc = CooldownCache(cache=router_cache, default_cooldown_time=60.0) + model_id = "busy-router-deployment" + + cc.add_deployment_to_cooldown( + model_id=model_id, + original_exception=Exception("Internal server error"), + exception_status=500, + cooldown_time=30.0, + ) + + for i in range(400): + router_cache.set_cache(key=f"unrelated-router-key-{i}", value={"n": i}) + + active = await cc.async_get_active_cooldowns([model_id], parent_otel_span=None) + assert [model_id] == [entry[0] for entry in active], ( + "unrelated router cache traffic must not evict a cooldown that is still running" + ) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index 4768988fc87..7ee0ed3701b 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -1,6 +1,8 @@ from unittest.mock import MagicMock, patch import litellm +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.router_utils.cooldown_handlers import ( _get_deployment_cooldown_policy, _resolve_allowed_fails_from_policy, @@ -269,18 +271,20 @@ class TestShouldCooldownBasedOnDeploymentPolicy: class TestShouldCooldownBasedOnAllowedFailsPolicy: - def _make_router(self, cooldown_time: float = 60.0) -> MagicMock: + def _make_router(self, cooldown_time: float = 60.0, cache: DualCache | None = None) -> MagicMock: router = MagicMock() router.cooldown_time = cooldown_time router.allowed_fails = 0 router.allowed_fails_policy = None router.get_allowed_fails_from_policy.return_value = None - router.failed_calls.get_cache.return_value = None + router.cache = cache if cache is not None else DualCache(in_memory_cache=InMemoryCache()) return router def test_cooldown_time_override_zero_is_not_falsy(self): """cooldown_time_override=0 must be honored; it must not fall through to the router-level value.""" router = self._make_router(cooldown_time=60.0) + router.cache = MagicMock() + router.cache.increment_cache.return_value = 1 exc = litellm.RateLimitError("429", "openai", "gpt-4") should_cooldown_based_on_allowed_fails_policy( @@ -291,12 +295,68 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy: cooldown_time_override=0.0, ) - set_cache_call = router.failed_calls.set_cache.call_args - assert set_cache_call is not None - assert set_cache_call[1]["ttl"] == 0.0, ( + increment_call = router.cache.increment_cache.call_args + assert increment_call is not None + assert increment_call[1]["ttl"] == 0.0, ( "cooldown_time_override=0 should be used as TTL, not the router-level 60.0" ) + def test_fail_counter_is_shared_across_router_instances(self): + """Two workers (two Router objects over one shared cache) must pool their failures toward allowed_fails.""" + shared_cache = DualCache(in_memory_cache=InMemoryCache()) + workers = (self._make_router(cache=shared_cache), self._make_router(cache=shared_cache)) + exc = litellm.AuthenticationError("401", "openai", "gpt-4") + + results = [ + should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=workers[i % 2], + deployment="dep-1", + original_exception=exc, + allowed_fails_override=5, + ) + for i in range(6) + ] + + assert results == [False, False, False, False, False, True] + assert shared_cache.get_cache(key="deployment:dep-1:allowed_fails") == 6 + + def test_fleet_wide_count_from_redis_decides_cooldown(self): + """The Redis (fleet-wide) count decides, even when this process has only seen one failure.""" + redis_cache = MagicMock() + redis_cache.increment_cache.return_value = 6 + router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)) + exc = litellm.AuthenticationError("401", "openai", "gpt-4") + + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="dep-1", + original_exception=exc, + allowed_fails_override=5, + ) + + assert result is True + redis_cache.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0) + + def test_redis_outage_falls_back_to_this_workers_count(self): + """When every Redis increment fails, the worker's own in-memory count must still cool the deployment down.""" + redis_cache = MagicMock() + redis_cache.increment_cache.side_effect = ConnectionError("redis down") + router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)) + exc = litellm.AuthenticationError("401", "openai", "gpt-4") + + results = [ + should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="dep-1", + original_exception=exc, + allowed_fails_override=5, + ) + for _ in range(6) + ] + + assert results == [False, False, False, False, False, True] + assert redis_cache.increment_cache.call_count == 6 + class TestRoutingGroupCooldownAlternatives: def _router(self, routing_groups=None): diff --git a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py index 6effbc5fa7f..9021d842daa 100644 --- a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py +++ b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py @@ -179,7 +179,7 @@ class TestHealthCheckCooldownIntegration: assert result is False # Check counter was incremented - current_fails = router.failed_calls.get_cache(key="deploy-1") + current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails") assert current_fails == 1 def test_health_check_failure_triggers_cooldown_at_threshold(self): @@ -263,7 +263,7 @@ class TestHealthCheckCooldownIntegration: assert "exception" not in healthy_endpoint # Verify failed_calls counter is untouched - current_fails = router.failed_calls.get_cache(key="deploy-1") + current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails") assert current_fails is None def test_disable_cooldowns_prevents_health_check_cooldown(self): 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_bedrock_extended_beta_models.py b/tests/test_litellm/test_bedrock_extended_beta_models.py deleted file mode 100644 index ebbbd6cab5c..00000000000 --- a/tests/test_litellm/test_bedrock_extended_beta_models.py +++ /dev/null @@ -1,170 +0,0 @@ -""" -Test suite for AWS Bedrock extended beta model support -Tests model configuration, pricing, and regional availability for: -- DeepSeek V3.2 -- Minimax M2.1 -- Moonshot AI Kimi K2.5 -- Qwen3 Coder Next -""" - -import os - -# Set env var to use local model cost map instead of fetching from remote -os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true" - -import pytest - -from litellm import get_model_info - -# Model configurations: (model_name, regions, max_input, max_output) -MODEL_CONFIGS = [ - ( - "deepseek.v3.2", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-north-1", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 163840, - 163840, - ), - ( - "minimax.minimax-m2.1", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-central-1", - "eu-north-1", - "eu-south-1", - "eu-west-1", - "eu-west-2", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 196000, - 8192, - ), - ( - "moonshotai.kimi-k2.5", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-north-1", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 262144, - 262144, - ), - ( - "qwen.qwen3-coder-next", - [ - "ap-northeast-1", - "ap-south-1", - "ap-southeast-3", - "eu-central-1", - "eu-south-1", - "eu-west-1", - "eu-west-2", - "sa-east-1", - "us-east-1", - "us-east-2", - "us-west-2", - ], - 262144, - 8192, - ), -] - - -class TestBedrockNewModels: - """Unified test suite for all new Bedrock models""" - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_model_info_primary_region( - self, model_name, regions, max_input, max_output - ): - """Test model configuration in primary region (us-east-1)""" - model = f"bedrock/us-east-1/{model_name}" - model_info = get_model_info(model) - - assert model_info is not None, f"Model {model_name} not found" - assert model_info["max_input_tokens"] == max_input - assert model_info["max_output_tokens"] == max_output - assert model_info["litellm_provider"] == "bedrock" - assert model_info["mode"] == "chat" - assert model_info["supports_function_calling"] is True - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_pricing_configured(self, model_name, regions, max_input, max_output): - """Verify pricing is set for all models""" - model = f"bedrock/us-east-1/{model_name}" - model_info = get_model_info(model) - - assert ( - model_info["input_cost_per_token"] > 0 - ), f"Missing input cost for {model_name}" - assert ( - model_info["output_cost_per_token"] > 0 - ), f"Missing output cost for {model_name}" - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_region_count(self, model_name, regions, max_input, max_output): - """Verify each bedrock/{region}/{model_name} resolves via get_model_info""" - for region in regions: - model = f"bedrock/{region}/{model_name}" - model_info = get_model_info(model) - assert model_info is not None, f"Model {model_name} not found in {region}" - assert model_info["max_input_tokens"] == max_input - assert model_info["max_output_tokens"] == max_output - - @pytest.mark.parametrize("model_name,regions,max_input,max_output", MODEL_CONFIGS) - def test_sample_regional_variants(self, model_name, regions, max_input, max_output): - """Test sample regional variants (us-east-1, eu-west-1, ap-northeast-1)""" - for region in ["us-east-1", "ap-northeast-1"]: - if region in regions: - model = f"bedrock/{region}/{model_name}" - model_info = get_model_info(model) - assert ( - model_info is not None - ), f"Model {model_name} not found in {region}" - assert model_info["max_input_tokens"] == max_input - assert model_info["litellm_provider"] == "bedrock" - - -class TestModelSpecificFeatures: - """Model-specific capability tests""" - - def test_deepseek_v3_2_context_window(self): - """DeepSeek V3.2 has 163K context window""" - model_info = get_model_info("bedrock/us-east-1/deepseek.v3.2") - assert model_info["max_input_tokens"] == 163840 - - def test_minimax_m2_1_context_window(self): - """Minimax M2.1 has 196K input, 8K output""" - model_info = get_model_info("bedrock/us-east-1/minimax.minimax-m2.1") - assert model_info["max_input_tokens"] == 196000 - assert model_info["max_output_tokens"] == 8192 - - def test_moonshotai_kimi_k2_5_context_window(self): - """Moonshot AI Kimi K2.5 has 256K context window""" - model_info = get_model_info("bedrock/us-east-1/moonshotai.kimi-k2.5") - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - - def test_qwen3_coder_next_context_window(self): - """Qwen3 Coder Next has 256K input, 8K output""" - model_info = get_model_info("bedrock/us-east-1/qwen.qwen3-coder-next") - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 8192 diff --git a/tests/test_litellm/test_bedrock_nemotron_super.py b/tests/test_litellm/test_bedrock_nemotron_super.py deleted file mode 100644 index 969db890e84..00000000000 --- a/tests/test_litellm/test_bedrock_nemotron_super.py +++ /dev/null @@ -1,51 +0,0 @@ -""" -Test suite for NVIDIA Nemotron Super 3 120B on AWS Bedrock -Verifies model configuration, pricing, and regional availability. -""" - -import os - -os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "true" - -import pytest - -from litellm import get_model_info - - -MODEL_NAME = "nvidia.nemotron-super-3-120b" - - -class TestNemotronSuper3120B: - """Test model definition for nvidia.nemotron-super-3-120b""" - - def test_model_info_primary_region(self): - """Test model resolves in us-east-1""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info is not None, f"Model {MODEL_NAME} not found" - assert model_info["max_input_tokens"] == 256000 - assert model_info["max_output_tokens"] == 32768 - assert model_info["litellm_provider"] == "bedrock_converse" - assert model_info["mode"] == "chat" - assert model_info["supports_function_calling"] is True - - def test_pricing_configured(self): - """Verify pricing matches AWS Bedrock rates""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info["input_cost_per_token"] == 1.5e-07 - assert model_info["output_cost_per_token"] == 6.5e-07 - - def test_context_window(self): - """Nemotron Super 3 120B has 256K input, 32K output on Bedrock""" - model_info = get_model_info(f"bedrock/us-east-1/{MODEL_NAME}") - - assert model_info["max_input_tokens"] == 256000 - assert model_info["max_output_tokens"] == 32768 - - def test_resolves_without_region(self): - """Test model resolves with just bedrock/ prefix""" - model_info = get_model_info(f"bedrock/{MODEL_NAME}") - - assert model_info is not None, f"Model {MODEL_NAME} not found without region" - assert model_info["max_input_tokens"] == 256000 diff --git a/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py b/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py deleted file mode 100644 index 1312aa110d3..00000000000 --- a/tests/test_litellm/test_bedrock_usgov_haiku_1hr_cache.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Validate that AWS GovCloud (Bedrock us-gov-*) Haiku 4.5 entries carry -the 1-hour cache write tier. - -AWS Bedrock GovCloud pricing applies a +20% premium over global -Anthropic rates. Global Haiku 4.5 1h cache write is $2.00/MTok; us-gov -is therefore $2.40/MTok — exactly 1.6x the 5-minute rate of $1.50/MTok. - -Source: https://aws.amazon.com/bedrock/pricing/ -""" - -import json -import os - -import pytest - - -@pytest.fixture(scope="module") -def model_data(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - return json.load(f) - - -HAIKU_USGOV_KEYS = [ - "bedrock/us-gov-east-1/anthropic.claude-haiku-4-5-20251001-v1:0", - "bedrock/us-gov-west-1/anthropic.claude-haiku-4-5-20251001-v1:0", -] - - -@pytest.mark.parametrize("model_key", HAIKU_USGOV_KEYS) -def test_usgov_haiku_4_5_1hr_cache_write(model_data, model_key): - assert model_key in model_data, f"Missing model entry: {model_key}" - info = model_data[model_key] - assert ( - info["cache_creation_input_token_cost"] == 1.5e-06 - ), f"{model_key}: 5m cache write should be $1.50/MTok" - assert ( - info["cache_creation_input_token_cost_above_1hr"] == 2.4e-06 - ), f"{model_key}: 1h cache write should be $2.40/MTok" - ratio = ( - info["cache_creation_input_token_cost_above_1hr"] - / info["cache_creation_input_token_cost"] - ) - assert abs(ratio - 1.6) < 1e-9, f"{model_key}: 1h/5m ratio is {ratio}, expected 1.6" diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py index 3576834dd27..3dfd7350a06 100644 --- a/tests/test_litellm/test_bedrock_usgov_pricing.py +++ b/tests/test_litellm/test_bedrock_usgov_pricing.py @@ -31,34 +31,6 @@ def model_data(): return json.load(f) -SONNET_4_5_USGOV_KEYS = [ - "bedrock/us-gov-east-1/anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-east-1/claude-sonnet-4-5-20250929-v1:0", - "bedrock/us-gov-west-1/claude-sonnet-4-5-20250929-v1:0", - "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0", -] - - -@pytest.mark.parametrize("model_key", SONNET_4_5_USGOV_KEYS) -def test_usgov_sonnet_4_5_pricing(model_data, model_key): - """Each us-gov sonnet-4-5 entry must carry the +20%-over-global rates - that AWS publishes on the GovCloud pricing page. - """ - assert model_key in model_data, f"Missing model entry: {model_key}" - info = model_data[model_key] - - assert info["input_cost_per_token"] == 3.6e-06, ( - f"{model_key}: input_cost_per_token should be $3.60/MTok (got {info['input_cost_per_token']})" - ) - assert info["output_cost_per_token"] == 1.8e-05, f"{model_key}: output_cost_per_token should be $18.00/MTok" - assert info["cache_creation_input_token_cost"] == 4.5e-06, f"{model_key}: 5m cache write should be $4.50/MTok" - assert info["cache_creation_input_token_cost_above_1hr"] == 7.2e-06, ( - f"{model_key}: 1h cache write should be $7.20/MTok" - ) - assert info["cache_read_input_token_cost"] == 3.6e-07, f"{model_key}: cache read should be $0.36/MTok" - - def test_usgov_carries_20_percent_premium_over_global(model_data): """The us-gov rates must equal 1.2x the global anthropic.* rates, matching AWS's documented GovCloud uplift. @@ -92,18 +64,6 @@ EXPECTED_USGOV_ABOVE_200K = { } -@pytest.mark.parametrize("field,expected", EXPECTED_USGOV_ABOVE_200K.items()) -def test_usgov_cross_region_above_200k_carries_gov_premium(model_data, field, expected): - """The `_above_200k_tokens` tier on the us-gov cross-region inference - profile must also carry the +20% GovCloud uplift. The original PR - corrected the base rates but left the 200k-tier fields at the +10% - commercial-US rates, undercharging long-context requests. - """ - info = model_data[USGOV_CROSS_REGION_KEY] - assert field in info, f"{USGOV_CROSS_REGION_KEY}: missing field {field}" - assert info[field] == expected, f"{USGOV_CROSS_REGION_KEY}: {field} should be {expected} (got {info[field]})" - - def test_usgov_cross_region_above_200k_ratio_to_global(model_data): """Cross-check via the property-based invariant: every `_above_200k_tokens` field on the us-gov cross-region profile must equal 1.2x the global @@ -117,167 +77,6 @@ def test_usgov_cross_region_above_200k_ratio_to_global(model_data): assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" -CLAUDE_GOV_EXPECTED = { - "anthropic.claude-sonnet-5": { - "input_cost_per_token": 2.4e-06, - "output_cost_per_token": 1.2e-05, - "cache_creation_input_token_cost": 3e-06, - "cache_creation_input_token_cost_above_1hr": 4.8e-06, - "cache_read_input_token_cost": 2.4e-07, - }, - "anthropic.claude-opus-4-8": { - "input_cost_per_token": 6e-06, - "output_cost_per_token": 3e-05, - "cache_creation_input_token_cost": 7.5e-06, - "cache_creation_input_token_cost_above_1hr": 1.2e-05, - "cache_read_input_token_cost": 6e-07, - }, - "anthropic.claude-opus-5": { - "input_cost_per_token": 6e-06, - "output_cost_per_token": 3e-05, - "cache_creation_input_token_cost": 7.5e-06, - "cache_creation_input_token_cost_above_1hr": 1.2e-05, - "cache_read_input_token_cost": 6e-07, - }, - "anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.2e-05, - "output_cost_per_token": 6e-05, - "cache_creation_input_token_cost": 1.5e-05, - "cache_creation_input_token_cost_above_1hr": 2.4e-05, - "cache_read_input_token_cost": 3e-07, - }, -} - - -USGOV_CLAUDE_KEY_TEMPLATES = { - "bedrock/us-gov-east-1/{base_key}": "bedrock", - "bedrock/us-gov-west-1/{base_key}": "bedrock", - "us-gov.{base_key}": "bedrock_converse", -} - - -@pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED) -@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items()) -def test_usgov_claude_pricing(model_data, key_template, expected_provider, base_key): - """Sonnet 5, Opus 4.8, Opus 5, and Fable 5.1 gov entries, both in-region keys - and the us-gov. geo inference profile the model cards list for GovCloud, must - carry the 1.2x GovCloud premium over the global anthropic.* rates. No public - AWS source (offer files, pricing page) lists Claude GovCloud rows; the premium - is the one AWS quotes for Opus 4.8 in GovCloud ($6/$30 per million). - """ - gov_key = key_template.format(base_key=base_key) - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == expected_provider - assert "search_context_cost_per_query" not in info - for field, expected in CLAUDE_GOV_EXPECTED[base_key].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - ratio = info[field] / model_data[base_key][field] - assert abs(ratio - 1.2) < 1e-9, f"{gov_key}: {field} gov/global ratio is {ratio}, expected 1.2" - - -CONVERSE_GOV_EXPECTED = { - "nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07), - "nvidia.nemotron-nano-9b-v2": (7.2e-08, 2.76e-07), - "nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07), - "nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07), - "openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07), - "openai.gpt-oss-120b-1:0": (1.8e-07, 7.2e-07), -} - - -@pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED) -@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items()) -def test_usgov_converse_model_pricing(model_data, key_template, expected_provider, base_key): - """Nemotron and gpt-oss gov entries, in-region and the us-gov. geo inference - profile both GovCloud regions list as ACTIVE, must match the AWS Bedrock - offer file, which prices both regions identically at 1.2x commercial. - """ - gov_key = key_template.format(base_key=base_key) - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key] - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["litellm_provider"] == expected_provider - base = model_data[base_key] - assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9 - assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9 - - -def test_usgov_west_llama3_8b_output_price_fixed(model_data): - """The us-gov-west-1 llama3-8b entry carried the 70B output rate ($2.65/MTok); - the AWS Bedrock offer file prices output at $0.60/MTok. AWS lists the model - in us-gov-west-1 only, so there is no east entry to check. - """ - info = model_data["bedrock/us-gov-west-1/meta.llama3-8b-instruct-v1:0"] - assert info["input_cost_per_token"] == 3e-07 - assert info["output_cost_per_token"] == 6e-07 - - -MANTLE_GOV_TIERED_EXPECTED = { - "openai.gpt-5.6-luna": { - "input_cost_per_token": 2.64e-07, - "input_cost_per_token_above_272k_tokens": 5.28e-07, - "cache_creation_input_token_cost": 3.3e-07, - "cache_creation_input_token_cost_above_272k_tokens": 6.6e-07, - "cache_read_input_token_cost": 2.64e-08, - "cache_read_input_token_cost_above_272k_tokens": 5.28e-08, - "output_cost_per_token": 1.584e-06, - "output_cost_per_token_above_272k_tokens": 2.376e-06, - }, - "openai.gpt-5.6-terra": { - "input_cost_per_token": 2.64e-06, - "input_cost_per_token_above_272k_tokens": 5.28e-06, - "cache_creation_input_token_cost": 3.3e-06, - "cache_creation_input_token_cost_above_272k_tokens": 6.6e-06, - "cache_read_input_token_cost": 2.64e-07, - "cache_read_input_token_cost_above_272k_tokens": 5.28e-07, - "output_cost_per_token": 1.584e-05, - "output_cost_per_token_above_272k_tokens": 2.376e-05, - }, -} - - -@pytest.mark.parametrize("model", MANTLE_GOV_TIERED_EXPECTED) -def test_usgov_west_mantle_terra_luna_pricing(model_data, model): - """Terra and Luna carry 1.2x commercial across every tier in the - us-gov-west-1 offer file; the us-gov-east-1 offer file has no SKUs for them. - """ - gov_key = f"bedrock_mantle/us-gov-west-1/{model}" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - for field, expected in MANTLE_GOV_TIERED_EXPECTED[model].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - assert info["litellm_provider"] == "bedrock_mantle" - assert f"bedrock_mantle/us-gov-east-1/{model}" not in model_data - - -@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"]) -def test_usgov_mantle_gpt_5_4_pricing_has_no_long_context_tier(model_data, region): - """gpt-5.4 gov rates come from the offer file, which publishes only the - standard tier in GovCloud: no long-context SKUs exist there, unlike commercial. - """ - gov_key = f"bedrock_mantle/{region}/openai.gpt-5.4" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["input_cost_per_token"] == 3.3e-06 - assert info["cache_read_input_token_cost"] == 3.3e-07 - assert info["output_cost_per_token"] == 1.98e-05 - assert not any(field.endswith("_above_272k_tokens") for field in info) - - -def test_usgov_mantle_grok_4_3_west_only(model_data): - """grok-4.3 is priced in the us-gov-west-1 offer file only; the east offer - file carries grok-4.6 instead. - """ - info = model_data["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 3e-06 - assert info["cache_read_input_token_cost"] == 2.4e-07 - assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data - - def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data): """us-gov-east-1 serves claude-3-haiku through the us-gov. inference profile only, so the profile row must bill exactly like the in-region gov row. @@ -290,96 +89,6 @@ def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data): } -GROK_4_6_GOV_KEYS = { - "us-gov.xai.grok-4.6": ("us.xai.grok-4.6", "bedrock_converse"), - "bedrock_mantle/us-gov-west-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"), - "bedrock_mantle/us-gov-east-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"), -} - - -@pytest.mark.parametrize("gov_key", GROK_4_6_GOV_KEYS) -def test_usgov_grok_4_6_pricing(model_data, gov_key): - """Both GovCloud regions serve grok-4.6 through the us-gov. profile only, and - both offer files price its standard SKU at 1.2x the commercial US rate. - """ - base_key, expected_provider = GROK_4_6_GOV_KEYS[gov_key] - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == expected_provider - assert info["input_cost_per_token"] == 2.64e-06 - assert info["output_cost_per_token"] == 7.92e-06 - assert info["cache_read_input_token_cost"] == 6.6e-07 - for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): - assert abs(info[field] / model_data[base_key][field] - 1.2) < 1e-9 - - -NOVA_GOV_WEST_EXPECTED = { - "amazon.nova-lite-v1:0": (7.2e-08, 2.88e-07), - "amazon.nova-micro-v1:0": (4.2e-08, 1.68e-07), -} - - -@pytest.mark.parametrize("base_key", NOVA_GOV_WEST_EXPECTED) -def test_usgov_west_nova_lite_micro_pricing(model_data, base_key): - """Nova Lite and Micro are on-demand in us-gov-west-1 only; the offer file - prices them at 1.2x commercial, like the Nova Pro row that was already there. - """ - gov_key = f"bedrock/us-gov-west-1/{base_key}" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - expected_input, expected_output = NOVA_GOV_WEST_EXPECTED[base_key] - assert info["litellm_provider"] == "bedrock" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert abs(info["input_cost_per_token"] / model_data[base_key]["input_cost_per_token"] - 1.2) < 1e-9 - assert abs(info["output_cost_per_token"] / model_data[base_key]["output_cost_per_token"] - 1.2) < 1e-9 - assert f"bedrock/us-gov-east-1/{base_key}" not in model_data - - -def test_usgov_west_nova_2_multimodal_embeddings_pricing(model_data): - """Every meter of the multimodal embedding model (tokens, images, audio and - video seconds) carries the 1.2x uplift the us-gov-west-1 offer file lists. - """ - gov_key = "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0" - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == "bedrock" - assert info["mode"] == "embedding" - assert info["input_cost_per_token"] == 1.62e-07 - assert info["input_cost_per_image"] == 7.2e-05 - assert info["input_cost_per_audio_per_second"] == 0.000168 - assert info["input_cost_per_video_per_second"] == 0.00084 - assert "bedrock/us-gov-east-1/amazon.nova-2-multimodal-embeddings-v1:0" not in model_data - - -MANTLE_GOV_FLAT_EXPECTED = { - "google.gemma-4-e2b": (4.8e-08, 9.6e-08, ("us-gov-west-1",)), - "google.gemma-4-26b-a4b": (1.56e-07, 4.8e-07, ("us-gov-west-1",)), - "google.gemma-4-31b": (1.68e-07, 4.8e-07, ("us-gov-west-1",)), - "openai.gpt-oss-20b": (8.4e-08, 3.6e-07, ("us-gov-west-1", "us-gov-east-1")), - "openai.gpt-oss-120b": (1.8e-07, 7.2e-07, ("us-gov-west-1", "us-gov-east-1")), -} - - -@pytest.mark.parametrize("model", MANTLE_GOV_FLAT_EXPECTED) -def test_usgov_mantle_gemma_and_gpt_oss_pricing(model_data, model): - """Gemma 4 is priced in the us-gov-west-1 offer file only and gpt-oss in both; - each Mantle gov row carries the offer file's standard SKU, and no row exists - for a region whose offer file has no SKU. - """ - expected_input, expected_output, regions = MANTLE_GOV_FLAT_EXPECTED[model] - for region in ("us-gov-west-1", "us-gov-east-1"): - gov_key = f"bedrock_mantle/{region}/{model}" - if region not in regions: - assert gov_key not in model_data - continue - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - assert info["litellm_provider"] == "bedrock_mantle" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - - GOV_ROW_SOURCES = { "us-gov.anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1", "bedrock/us-gov-west-1/anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1", @@ -417,33 +126,3 @@ def test_usgov_rows_keep_commercial_limits_and_capabilities(model_data, gov_key) assert _non_pricing_fields(gov) == _non_pricing_fields(model_data[GOV_ROW_SOURCES[gov_key]]) assert "search_context_cost_per_query" not in gov assert "source" not in gov - - -AZURE_GOV_EXPECTED = { - "azure/us-gov/gpt-5.1": { - "input_cost_per_token": 1.71875e-06, - "cache_read_input_token_cost": 1.71875e-07, - "output_cost_per_token": 1.375e-05, - }, - "azure/us-gov/o3-mini": { - "input_cost_per_token": 1.513e-06, - "cache_read_input_token_cost": 7.57e-07, - "output_cost_per_token": 6.05e-06, - }, - "azure/us-gov/text-embedding-3-large": {"input_cost_per_token": 1.63e-07}, - "azure/us-gov/text-embedding-3-small": {"input_cost_per_token": 2.5e-08}, -} - - -@pytest.mark.parametrize("gov_key", AZURE_GOV_EXPECTED) -def test_azure_usgov_pricing(model_data, gov_key): - """Azure Government meters from the Azure retail prices API - (usgovvirginia/usgovarizona, serviceName 'Foundry Models'). No Government - retirement schedule is published, so these entries carry no deprecation_date. - """ - assert gov_key in model_data, f"Missing model entry: {gov_key}" - info = model_data[gov_key] - for field, expected in AZURE_GOV_EXPECTED[gov_key].items(): - assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})" - assert info["litellm_provider"] == "azure" - assert "deprecation_date" not in info diff --git a/tests/test_litellm/test_claude_fable_5_config.py b/tests/test_litellm/test_claude_fable_5_config.py index 3ecf94602d9..0473161faac 100644 --- a/tests/test_litellm/test_claude_fable_5_config.py +++ b/tests/test_litellm/test_claude_fable_5_config.py @@ -14,7 +14,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -27,89 +26,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_fable_5_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = [ - ("claude-fable-5", "anthropic"), - ("anthropic.claude-fable-5", "bedrock_converse"), - ("vertex_ai/claude-fable-5", "vertex_ai-anthropic_models"), - # Unlike Opus 4.8 (200k on Foundry), Fable 5 has the full 1M context - # window on Microsoft Foundry. - ("azure_ai/claude-fable-5", "azure_ai"), - ] - - for model_name, provider in expected_models: - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # $10 / $50 per MTok (2x Opus 4.8), with the standard 1.25x 5m - # cache-write, 2x 1h cache-write, and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 1e-05 - assert info["output_cost_per_token"] == 5e-05 - assert info["cache_creation_input_token_cost"] == 1.25e-05 - assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 - assert info["cache_read_input_token_cost"] == 1e-06 - - # Flat-rate across the full 1M context window. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - - -def test_fable_5_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - # Fable 5 launched with us/eu geo inference profiles plus a global profile - # (no au/apac/jp). Global uses base pricing; geo profiles carry the - # standard 10% regional premium. - expected_models = { - "global.anthropic.claude-fable-5": { - "input_cost_per_token": 1e-05, - "output_cost_per_token": 5e-05, - "cache_creation_input_token_cost": 1.25e-05, - "cache_read_input_token_cost": 1e-06, - }, - "us.anthropic.claude-fable-5": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 1.1e-06, - }, - "eu.anthropic.claude-fable-5": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 1.1e-06, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - def test_fable_5_geo_multiplier_without_fast_mode(): """First-party ``inference_geo='us'`` carries the 1.1x premium, but unlike the Opus line there is no fast-mode variant for Fable 5; a ``fast`` key @@ -144,13 +60,6 @@ def test_fable_5_registered_for_bedrock_converse(): assert "anthropic.claude-fable-5" in BEDROCK_CONVERSE_MODELS -def test_fable_5_provider_resolves_via_model_info(local_model_cost_map): - info = litellm.get_model_info(model="claude-fable-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], @@ -222,46 +131,6 @@ FABLE_5_1_VARIANTS = ( ) -def test_fable_5_1_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = [ - ("claude-fable-5-1", "anthropic"), - ("anthropic.claude-fable-5-1", "bedrock_converse"), - ("vertex_ai/claude-fable-5-1", "vertex_ai-anthropic_models"), - ("azure_ai/claude-fable-5-1", "azure_ai"), - ] - - for model_name, provider in expected_models: - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["input_cost_per_token"] == 1e-05 - assert info["output_cost_per_token"] == 5e-05 - assert info["cache_creation_input_token_cost"] == 1.25e-05 - assert info["cache_creation_input_token_cost_above_1hr"] == 2e-05 - - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_forced_tool_use"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - assert info["prompt_cache_min_tokens"] == 512 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], @@ -280,48 +149,6 @@ def test_fable_5_1_cache_reads_cost_a_quarter_of_fable_5(cost_map): ), model_name -def test_fable_5_1_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - expected_models = { - "global.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1e-05, - "output_cost_per_token": 5e-05, - "cache_creation_input_token_cost": 1.25e-05, - "cache_read_input_token_cost": 2.5e-07, - }, - "us.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 2.75e-07, - }, - "eu.anthropic.claude-fable-5-1": { - "input_cost_per_token": 1.1e-05, - "output_cost_per_token": 5.5e-05, - "cache_creation_input_token_cost": 1.375e-05, - "cache_read_input_token_cost": 2.75e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - -def test_fable_5_1_geo_multiplier_without_fast_mode(): - """Fable 5.1 has no fast mode, so a ``fast`` key here would misprice - ``speed='fast'`` requests.""" - model_data = _load_root_cost_map() - assert model_data["claude-fable-5-1"]["provider_specific_entry"] == {"us": 1.1} - - def test_fable_5_1_present_in_bundled_backup(): backup = GetModelCostMap.load_local_model_cost_map() root = _load_root_cost_map() @@ -334,13 +161,6 @@ def test_fable_5_1_registered_for_bedrock_converse(): assert "anthropic.claude-fable-5-1" in BEDROCK_CONVERSE_MODELS -def test_fable_5_1_provider_resolves_via_model_info(local_model_cost_map): - info = litellm.get_model_info(model="claude-fable-5-1") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "model", [ diff --git a/tests/test_litellm/test_claude_haiku_4_5_config.py b/tests/test_litellm/test_claude_haiku_4_5_config.py index 8755e5d156f..9172b6479a5 100644 --- a/tests/test_litellm/test_claude_haiku_4_5_config.py +++ b/tests/test_litellm/test_claude_haiku_4_5_config.py @@ -7,57 +7,6 @@ import json import os -def test_bedrock_haiku_4_5_configuration(): - """Test that all Bedrock Claude Haiku 4.5 models use bedrock_converse provider""" - # Load model configuration - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - # All Bedrock Haiku 4.5 variants that should use bedrock_converse - bedrock_haiku_models = [ - "anthropic.claude-haiku-4-5-20251001-v1:0", - "anthropic.claude-haiku-4-5@20251001", - "us.anthropic.claude-haiku-4-5-20251001-v1:0", - "eu.anthropic.claude-haiku-4-5-20251001-v1:0", - "apac.anthropic.claude-haiku-4-5-20251001-v1:0", - "jp.anthropic.claude-haiku-4-5-20251001-v1:0", - "global.anthropic.claude-haiku-4-5-20251001-v1:0", - "au.anthropic.claude-haiku-4-5-20251001-v1:0", - ] - - for model in bedrock_haiku_models: - assert model in model_data, f"Model {model} not found in config" - model_info = model_data[model] - - # Verify uses bedrock_converse (not legacy bedrock provider) - assert ( - model_info["litellm_provider"] == "bedrock_converse" - ), f"{model} should use bedrock_converse provider, got {model_info['litellm_provider']}" - - # Verify supports vision (key missing capability) - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - # Verify core capabilities - assert model_info.get("supports_computer_use") is True - assert model_info.get("supports_function_calling") is True - assert model_info.get("supports_tool_choice") is True - assert model_info.get("supports_prompt_caching") is True - assert model_info.get("supports_response_schema") is True - assert model_info.get("supports_pdf_input") is True - assert model_info.get("supports_assistant_prefill") is True - assert model_info.get("supports_reasoning") is True - - # Verify token limits - assert model_info["max_input_tokens"] == 200000 - assert model_info["max_output_tokens"] == 64000 - assert model_info["mode"] == "chat" - - def test_bedrock_haiku_4_5_matches_sonnet_capabilities(): """ Test that Haiku 4.5 has same capabilities as Sonnet 4.5 @@ -97,36 +46,3 @@ def test_bedrock_haiku_4_5_matches_sonnet_capabilities(): assert haiku_info.get(capability) == sonnet_info.get( capability ), f"Capability {capability} mismatch: Haiku={haiku_info.get(capability)}, Sonnet={sonnet_info.get(capability)}" - - -def test_anthropic_api_haiku_4_5_configuration(): - """Test that Anthropic API Claude Haiku 4.5 has correct configuration""" - # Load model configuration - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - # Anthropic API models (not Bedrock) - anthropic_models = [ - "claude-haiku-4-5-20251001", - "claude-haiku-4-5", - ] - - for model in anthropic_models: - assert model in model_data, f"Model {model} not found in config" - model_info = model_data[model] - - # Should use anthropic provider (not bedrock) - assert ( - model_info["litellm_provider"] == "anthropic" - ), f"{model} should use anthropic provider" - - # Should support vision - assert ( - model_info.get("supports_vision") is True - ), f"{model} should support vision" - - # Should have larger output token limit (64K for Anthropic API) - assert model_info["max_output_tokens"] == 64000 diff --git a/tests/test_litellm/test_claude_opus_4_6_config.py b/tests/test_litellm/test_claude_opus_4_6_config.py index 89d2cd916e0..9a8632924f2 100644 --- a/tests/test_litellm/test_claude_opus_4_6_config.py +++ b/tests/test_litellm/test_claude_opus_4_6_config.py @@ -71,125 +71,6 @@ def test_claude_4_6_australia_region_uses_au_prefix_not_apac(): ), "apac.anthropic.claude-sonnet-4-6 should not be in bedrock_converse_models" -def test_opus_4_6_model_pricing_and_capabilities(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - expected_models = { - "claude-opus-4-6": { - "provider": "anthropic", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "claude-opus-4-6-20260205": { - "provider": "anthropic", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "anthropic.claude-opus-4-6-v1": { - "provider": "bedrock_converse", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "vertex_ai/claude-opus-4-6": { - "provider": "vertex_ai-anthropic_models", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - "azure_ai/claude-opus-4-6": { - "provider": "azure_ai", - "has_long_context_pricing": False, - "max_input_tokens": 1000000, - }, - } - - for model_name, config in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == config["provider"] - assert info["mode"] == "chat" - assert info["max_input_tokens"] == config["max_input_tokens"] - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_read_input_token_cost"] == 5e-07 - - if config["has_long_context_pricing"]: - assert info["input_cost_per_token_above_200k_tokens"] == 1e-05 - assert info["output_cost_per_token_above_200k_tokens"] == 3.75e-05 - assert info["cache_creation_input_token_cost_above_200k_tokens"] == 1.25e-05 - assert info["cache_read_input_token_cost_above_200k_tokens"] == 1e-06 - else: - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - assert "cache_creation_input_token_cost_above_200k_tokens" not in info - assert "cache_read_input_token_cost_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_opus_4_6_bedrock_regional_model_pricing(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - expected_models = { - "global.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_read_input_token_cost": 5e-07, - }, - "us.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "eu.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "au.anthropic.claude-opus-4-6-v1": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - assert info["supports_assistant_prefill"] is False - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - assert "cache_creation_input_token_cost_above_200k_tokens" not in info - assert "cache_read_input_token_cost_above_200k_tokens" not in info - for key, value in expected.items(): - assert info[key] == value - - def test_opus_4_6_alias_and_dated_metadata_match(): json_path = os.path.join( os.path.dirname(__file__), "../../model_prices_and_context_window.json" diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/test_litellm/test_claude_opus_4_8_config.py index 760512ad31b..e75fdba54ed 100644 --- a/tests/test_litellm/test_claude_opus_4_8_config.py +++ b/tests/test_litellm/test_claude_opus_4_8_config.py @@ -16,7 +16,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -29,102 +28,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_opus_4_8_model_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_models = { - "claude-opus-4-8": { - "provider": "anthropic", - "max_input_tokens": 1000000, - }, - "anthropic.claude-opus-4-8": { - "provider": "bedrock_converse", - "max_input_tokens": 1000000, - }, - "vertex_ai/claude-opus-4-8": { - "provider": "vertex_ai-anthropic_models", - "max_input_tokens": 1000000, - }, - "azure_ai/claude-opus-4-8": { - "provider": "azure_ai", - "max_input_tokens": 1000000, - }, - } - - for model_name, config in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == config["provider"] - assert info["mode"] == "chat" - assert info["max_input_tokens"] == config["max_input_tokens"] - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Base pricing matches Opus 4.7: $5 / $25 per MTok, with the standard - # 1.25x cache-write and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_read_input_token_cost"] == 5e-07 - - # Opus 4.x flagships are flat-rate across the full context window. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - assert info["supports_assistant_prefill"] is False - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - assert model_data["claude-opus-4-8"]["supports_native_structured_output"] is True - - -def test_opus_4_8_bedrock_regional_model_pricing(): - model_data = _load_root_cost_map() - - # Global endpoints use base pricing; regional endpoints carry a 10% premium. - expected_models = { - "global.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_read_input_token_cost": 5e-07, - }, - "us.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "eu.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - "au.anthropic.claude-opus-4-8": { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_read_input_token_cost": 5.5e-07, - }, - } - - for model_name, expected in expected_models.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in expected.items(): - assert info[key] == value - - def test_opus_4_8_fast_mode_multiplier(): """Opus 4.8 dropped fast-mode pricing to 2x base ($10/$50 per MTok); Opus 4.7 was 6x ($30/$150).""" @@ -134,44 +37,10 @@ def test_opus_4_8_fast_mode_multiplier(): assert entry["fast"] == 2.0 -def test_opus_4_8_present_in_bundled_backup(): - """The bundled backup is the runtime fallback (and what tests load with - ``LITELLM_LOCAL_MODEL_COST_MAP=True``) — it must carry the same entries as - the root cost map, otherwise the model resolves on one path but not the - other.""" - backup = GetModelCostMap.load_local_model_cost_map() - for model_name in ( - "claude-opus-4-8", - "anthropic.claude-opus-4-8", - "global.anthropic.claude-opus-4-8", - "us.anthropic.claude-opus-4-8", - "eu.anthropic.claude-opus-4-8", - "au.anthropic.claude-opus-4-8", - "vertex_ai/claude-opus-4-8", - "vertex_ai/claude-opus-4-8@default", - "azure_ai/claude-opus-4-8", - ): - assert model_name in backup, f"Missing from backup cost map: {model_name}" - assert backup["claude-opus-4-8"]["supports_native_structured_output"] is True - - def test_opus_4_8_registered_for_bedrock_converse(): assert "anthropic.claude-opus-4-8" in BEDROCK_CONVERSE_MODELS -def test_opus_4_8_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-opus-4-8`` must resolve to provider ``anthropic``. - - Before the cost-map entry existed, the model was unknown to LiteLLM, so it - could not be tied to the ``anthropic`` provider and an ``anthropic/*`` - wildcard deployment would not match it. - """ - info = litellm.get_model_info(model="claude-opus-4-8") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/test_litellm/test_claude_opus_5_config.py index 34744aad17b..285d556ef2b 100644 --- a/tests/test_litellm/test_claude_opus_5_config.py +++ b/tests/test_litellm/test_claude_opus_5_config.py @@ -17,7 +17,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -52,91 +51,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_opus_5_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_providers = { - "claude-opus-5": "anthropic", - "anthropic.claude-opus-5": "bedrock_converse", - "vertex_ai/claude-opus-5": "vertex_ai-anthropic_models", - "azure_ai/claude-opus-5": "azure_ai", - } - - for model_name, provider in expected_providers.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Opus 5 ships at Opus 4.8's rates: $5 / $25 per MTok, with the standard - # 1.25x cache-write, 2x 1-hour cache-write, and 0.1x cache-read multipliers. - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 2.5e-05 - assert info["cache_creation_input_token_cost"] == 6.25e-06 - assert info["cache_creation_input_token_cost_above_1hr"] == 1e-05 - assert info["cache_read_input_token_cost"] == 5e-07 - - # Flat rate across the full 1M window, no long-context premium. - assert "input_cost_per_token_above_200k_tokens" not in info - assert "output_cost_per_token_above_200k_tokens" not in info - - # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no - # assistant prefill. - assert info["supports_adaptive_thinking"] is True - assert info["supports_reasoning"] is True - assert info["supports_sampling_params"] is False - assert info["supports_assistant_prefill"] is False - assert info["supports_xhigh_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_opus_5_bedrock_regional_pricing(): - """Global/base endpoints use base pricing; the us./eu./au./jp. regional - cross-region inference profiles carry a 10% premium.""" - model_data = _load_root_cost_map() - - base_pricing = { - "input_cost_per_token": 5e-06, - "output_cost_per_token": 2.5e-05, - "cache_creation_input_token_cost": 6.25e-06, - "cache_creation_input_token_cost_above_1hr": 1e-05, - "cache_read_input_token_cost": 5e-07, - } - regional_pricing = { - "input_cost_per_token": 5.5e-06, - "output_cost_per_token": 2.75e-05, - "cache_creation_input_token_cost": 6.875e-06, - "cache_creation_input_token_cost_above_1hr": 1.1e-05, - "cache_read_input_token_cost": 5.5e-07, - } - - expected = { - "anthropic.claude-opus-5": base_pricing, - "global.anthropic.claude-opus-5": base_pricing, - "us.anthropic.claude-opus-5": regional_pricing, - "eu.anthropic.claude-opus-5": regional_pricing, - "au.anthropic.claude-opus-5": regional_pricing, - "jp.anthropic.claude-opus-5": regional_pricing, - } - - for model_name, pricing in expected.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - for key, value in pricing.items(): - assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}" - - @pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name): """Bedrock accepts every effort level for Opus 5, so no clamp belongs here. @@ -216,18 +130,6 @@ def test_opus_5_registered_for_bedrock_converse(): assert "anthropic.claude-opus-5" in BEDROCK_CONVERSE_MODELS -def test_opus_5_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-opus-5`` must resolve to provider ``anthropic``. - - Without the cost-map entry the model is unknown to LiteLLM, so it cannot be - tied to the ``anthropic`` provider and an ``anthropic/*`` wildcard deployment - would not match it.""" - info = litellm.get_model_info(model="claude-opus-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_claude_sonnet_5_config.py b/tests/test_litellm/test_claude_sonnet_5_config.py index 8504326cd21..8c6d2cd1851 100644 --- a/tests/test_litellm/test_claude_sonnet_5_config.py +++ b/tests/test_litellm/test_claude_sonnet_5_config.py @@ -15,7 +15,6 @@ import os import pytest -import litellm from litellm.constants import BEDROCK_CONVERSE_MODELS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap @@ -41,96 +40,6 @@ def _load_root_cost_map() -> dict: return json.load(f) - -def test_sonnet_5_pricing_and_capabilities(): - model_data = _load_root_cost_map() - - expected_providers = { - "claude-sonnet-5": "anthropic", - "anthropic.claude-sonnet-5": "bedrock_converse", - "vertex_ai/claude-sonnet-5": "vertex_ai-anthropic_models", - "azure_ai/claude-sonnet-5": "azure_ai", - } - - for model_name, provider in expected_providers.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - - assert info["litellm_provider"] == provider - assert info["mode"] == "chat" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - # Introductory Sonnet 5 pricing through 2026-08-31: $2 / $10 per MTok, - # with the 1.25x cache-write and 0.1x cache-read multipliers. On - # 2026-09-01 flip these five fields back to the sticker rate, here and - # in both cost-map JSON files (all ten claude-sonnet-5 entries): - # input_cost_per_token: 3e-06 - # output_cost_per_token: 1.5e-05 - # cache_creation_input_token_cost: 3.75e-06 - # cache_creation_input_token_cost_above_1hr: 6e-06 - # cache_read_input_token_cost: 3e-07 - # Regional Bedrock profiles (us./eu./au./jp.) stay at 1.1x those values: - # 3.3e-06 / 1.65e-05 / 4.125e-06 / 6.6e-06 / 3.3e-07 (see - # test_sonnet_5_bedrock_regional_pricing below). - assert info["input_cost_per_token"] == 2e-06 - assert info["output_cost_per_token"] == 1e-05 - assert info["cache_creation_input_token_cost"] == 2.5e-06 - assert info["cache_creation_input_token_cost_above_1hr"] == 4e-06 - assert info["cache_read_input_token_cost"] == 2e-07 - - # gen-5 adaptive-thinking profile: effort-driven, no sampling params, no - # assistant prefill. - assert info["supports_adaptive_thinking"] is True - assert info["supports_reasoning"] is True - assert info["supports_sampling_params"] is False - assert info["supports_assistant_prefill"] is False - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - - -def test_sonnet_5_bedrock_regional_pricing(): - """Global/base endpoints use base pricing; the us./eu./au./jp. regional - cross-region inference profiles carry a 10% premium.""" - model_data = _load_root_cost_map() - - base_pricing = { - "input_cost_per_token": 2e-06, - "output_cost_per_token": 1e-05, - "cache_creation_input_token_cost": 2.5e-06, - "cache_creation_input_token_cost_above_1hr": 4e-06, - "cache_read_input_token_cost": 2e-07, - } - regional_pricing = { - "input_cost_per_token": 2.2e-06, - "output_cost_per_token": 1.1e-05, - "cache_creation_input_token_cost": 2.75e-06, - "cache_creation_input_token_cost_above_1hr": 4.4e-06, - "cache_read_input_token_cost": 2.2e-07, - } - - expected = { - "anthropic.claude-sonnet-5": base_pricing, - "global.anthropic.claude-sonnet-5": base_pricing, - "us.anthropic.claude-sonnet-5": regional_pricing, - "eu.anthropic.claude-sonnet-5": regional_pricing, - "au.anthropic.claude-sonnet-5": regional_pricing, - "jp.anthropic.claude-sonnet-5": regional_pricing, - } - - for model_name, pricing in expected.items(): - assert model_name in model_data, f"Missing model entry: {model_name}" - info = model_data[model_name] - assert info["litellm_provider"] == "bedrock_converse" - assert info["bedrock_output_config_effort_ceiling"] == "xhigh" - for key, value in pricing.items(): - assert info[key] == value, f"{model_name}.{key} = {info[key]}, want {value}" - - def test_sonnet_5_present_in_bundled_backup(): """The bundled backup is the runtime fallback (and what tests load with ``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the @@ -144,18 +53,6 @@ def test_sonnet_5_registered_for_bedrock_converse(): assert "anthropic.claude-sonnet-5" in BEDROCK_CONVERSE_MODELS -def test_sonnet_5_provider_resolves_via_model_info(local_model_cost_map): - """Regression: ``claude-sonnet-5`` must resolve to provider ``anthropic``. - - Before the cost-map entry existed, the model was unknown to LiteLLM, so it - could not be tied to the ``anthropic`` provider and an ``anthropic/*`` - wildcard deployment would not match it.""" - info = litellm.get_model_info(model="claude-sonnet-5") - assert info["litellm_provider"] == "anthropic" - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "cost_map", [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], diff --git a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py index e33bcfb8378..4e770be7c3e 100644 --- a/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py +++ b/tests/test_litellm/test_cloudflare_workers_ai_model_metadata.py @@ -27,17 +27,6 @@ BACKUP_MAP = os.path.join( ) -@pytest.fixture(autouse=True) -def _use_local_model_cost_map(monkeypatch): - original_model_cost = litellm.model_cost - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - litellm.model_cost = litellm.get_model_cost_map(url="") - try: - yield - finally: - litellm.model_cost = original_model_cost - - def _load(path: str) -> dict: with open(path, encoding="utf-8") as f: return json.load(f) @@ -47,50 +36,6 @@ def _cloudflare_keys(data: dict) -> set: return {k for k in data if k.startswith("cloudflare/")} -def test_glm_5_2_entry_is_present_and_well_formed(): - entry = litellm.model_cost["cloudflare/@cf/zai-org/glm-5.2"] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["input_cost_per_token"] > 0 - assert entry["output_cost_per_token"] > 0 - - -def test_vision_model_is_flagged_supports_vision(): - entry = litellm.model_cost["cloudflare/@cf/meta/llama-3.2-11b-vision-instruct"] - assert entry["litellm_provider"] == "cloudflare" - assert entry.get("supports_vision") is True - - -def test_additional_current_models_are_present(): - for key in ( - "cloudflare/@cf/openai/gpt-oss-120b", - "cloudflare/@cf/meta/llama-3.3-70b-instruct-fp8-fast", - ): - entry = litellm.model_cost[key] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["input_cost_per_token"] > 0 - assert entry["output_cost_per_token"] > 0 - - -@pytest.mark.parametrize( - "key, published_price_per_audio_minute", - [ - ("cloudflare/@cf/openai/whisper", 0.00045), - ("cloudflare/@cf/openai/whisper-large-v3-turbo", 0.00051), - ], -) -def test_whisper_transcription_pricing_is_stored_per_second(key, published_price_per_audio_minute): - entry = litellm.model_cost[key] - assert entry["litellm_provider"] == "cloudflare" - assert entry["mode"] == "audio_transcription" - assert entry["supported_endpoints"] == ["/v1/audio/transcriptions"] - assert entry["output_cost_per_second"] == 0.0 - assert entry["input_cost_per_second"] == pytest.approx(published_price_per_audio_minute / 60) - - def test_root_and_backup_have_identical_cloudflare_keys(): if not os.path.exists(ROOT_MAP): pytest.skip("root cost map only ships in source checkouts") diff --git a/tests/test_litellm/test_component_entrypoint.py b/tests/test_litellm/test_component_entrypoint.py index b0969e2c694..09837d2b233 100644 --- a/tests/test_litellm/test_component_entrypoint.py +++ b/tests/test_litellm/test_component_entrypoint.py @@ -223,6 +223,59 @@ def test_gating_matches_the_monolithic_entrypoint_and_get_secret_bool( assert monolith[1] == ("args=litellm --port 4000" if traced else "args=--port 4000") +def test_wipes_the_prometheus_multiproc_dir_before_uvicorn_forks(tmp_path: Path) -> None: + """A restarted container inherits the emptyDir of its predecessor, whose worker pids it may reuse, so the + stale .db files must be gone before any worker opens the one carrying its own pid.""" + multiproc_dir = tmp_path / "multiproc" + multiproc_dir.mkdir() + (multiproc_dir / "gauge_livesum_7.db").write_bytes(b"stale") + (multiproc_dir / "counter_7.db").write_bytes(b"stale") + (multiproc_dir / "keep.txt").write_text("not a sample") + + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + _write_stubs(bin_dir, ("uvicorn",)) + record = tmp_path / "record.txt" + env = { + **os.environ, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "RECORD": str(record), + "PROMETHEUS_MULTIPROC_DIR": str(multiproc_dir), + } + env.pop("USE_DDTRACE", None) + result = subprocess.run( + ["sh", str(COMPONENT_ENTRYPOINT), "uvicorn", "gateway.main:app"], + env=env, + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}" + assert sorted(p.name for p in multiproc_dir.iterdir()) == ["keep.txt"] + assert record.read_text().splitlines()[0] == "exec=uvicorn" + + +def test_creates_a_missing_prometheus_multiproc_dir(tmp_path: Path) -> None: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + _write_stubs(bin_dir, ("uvicorn",)) + missing = tmp_path / "multiproc" + env = { + **os.environ, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "RECORD": str(tmp_path / "record.txt"), + "PROMETHEUS_MULTIPROC_DIR": str(missing), + } + env.pop("USE_DDTRACE", None) + result = subprocess.run( + ["sh", str(COMPONENT_ENTRYPOINT), "uvicorn", "gateway.main:app"], env=env, capture_output=True, text=True + ) + + assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}" + assert missing.is_dir() + + def _copied_script(dockerfile: Path, image_path: str) -> Path: """Resolve the repo file a Dockerfile `COPY`s to `image_path`, so tests run what the image ships.""" matches = _COPY_RE.findall(dockerfile.read_text()) 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_daybreak_model_metadata.py b/tests/test_litellm/test_daybreak_model_metadata.py index dbb7ecdffac..c3bac14dbbd 100644 --- a/tests/test_litellm/test_daybreak_model_metadata.py +++ b/tests/test_litellm/test_daybreak_model_metadata.py @@ -32,22 +32,6 @@ def _load(path): return json.load(f) -@pytest.mark.parametrize("model", DAYBREAK_MODELS) -def test_daybreak_capability_contract(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "openai" - assert info["mode"] == "chat" - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses"] - - assert info["supports_computer_use"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - - def test_blue_alias_matches_its_snapshot_computer_use(): cost_map = _load(MAIN_PATH) diff --git a/tests/test_litellm/test_deepseek_model_metadata.py b/tests/test_litellm/test_deepseek_model_metadata.py index b9eb33f0972..9cbd14ebd1e 100644 --- a/tests/test_litellm/test_deepseek_model_metadata.py +++ b/tests/test_litellm/test_deepseek_model_metadata.py @@ -39,26 +39,6 @@ class TestDeepSeekModelCostEntries: """Verify that provider-prefixed DeepSeek entries contain the same capability flags as their bare-name counterparts in the JSON files.""" - def test_deepseek_chat_supports_response_schema_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_response_schema") is True - - def test_deepseek_reasoner_supports_response_schema_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_response_schema") is True - - def test_deepseek_chat_supports_system_messages_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_system_messages") is True - - def test_deepseek_reasoner_supports_system_messages_in_backup(self): - data = _load_backup_json() - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_system_messages") is True - def test_deepseek_chat_max_input_tokens_matches_bare_in_backup(self): data = _load_backup_json() bare = data.get("deepseek-chat", {}) @@ -71,26 +51,6 @@ class TestDeepSeekModelCostEntries: prefixed = data.get("deepseek/deepseek-reasoner", {}) assert prefixed.get("max_output_tokens") == bare.get("max_output_tokens") - def test_main_json_deepseek_chat_supports_response_schema(self): - main_path = os.path.join( - os.path.dirname(os.path.dirname(litellm.__file__)), - "model_prices_and_context_window.json", - ) - with open(main_path, encoding="utf-8") as f: - data = json.load(f) - entry = data.get("deepseek/deepseek-chat", {}) - assert entry.get("supports_response_schema") is True - - def test_main_json_deepseek_reasoner_supports_response_schema(self): - main_path = os.path.join( - os.path.dirname(os.path.dirname(litellm.__file__)), - "model_prices_and_context_window.json", - ) - with open(main_path, encoding="utf-8") as f: - data = json.load(f) - entry = data.get("deepseek/deepseek-reasoner", {}) - assert entry.get("supports_response_schema") is True - # --------------------------------------------------------------------------- # API-level tests – verify supports_response_schema returns True diff --git a/tests/test_litellm/test_default_branch.py b/tests/test_litellm/test_default_branch.py new file mode 100644 index 00000000000..a1b2a8c5c91 --- /dev/null +++ b/tests/test_litellm/test_default_branch.py @@ -0,0 +1,240 @@ +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Final + +import pytest + +ROOT: Final = Path(__file__).resolve().parents[2] + + +def _git(repo: Path, *args: str) -> str: + return subprocess.run(["git", *args], cwd=repo, check=True, capture_output=True, text=True).stdout.strip() + + +def _commit(repo: Path, message: str) -> None: + _git(repo, "add", ".") + _git(repo, "-c", "user.name=Test", "-c", "user.email=test@example.com", "commit", "-qm", message) + + +@pytest.fixture +def remote_and_clone(tmp_path: Path) -> tuple[Path, Path]: + seed: Final = tmp_path / "seed" + seed.mkdir() + _git(seed, "init", "-q", "-b", "litellm_internal_staging") + (seed / "scripts").mkdir() + for name in ( + "default_branch.py", + "budget_ratchet_check.py", + "ruff_strict_gate.py", + "type_discipline_gate.py", + "test_quality_gate.py", + "type_check_gate.py", + "gate_slot_lock.py", + ): + shutil.copyfile(ROOT / "scripts" / name, seed / "scripts" / name) + shutil.copyfile(ROOT / "Makefile", seed / "Makefile") + (seed / "litellm").mkdir() + (seed / "litellm" / "example.py").write_text("value = 0\n") + (seed / "ruff-strict-budget.json").write_text('{"C901": {"limit": 1}}\n') + _commit(seed, "staging base") + _git(seed, "checkout", "-qb", "main") + (seed / "litellm" / "example.py").write_text("value = 1\n") + (seed / "ruff-strict-budget.json").write_text('{"C901": {"limit": 0}}\n') + _commit(seed, "main base") + remote: Final = tmp_path / "remote.git" + _git(tmp_path, "clone", "-q", "--bare", str(seed), str(remote)) + _git(remote, "symbolic-ref", "HEAD", "refs/heads/litellm_internal_staging") + repo: Final = tmp_path / "clone" + _git(tmp_path, "clone", "-q", "--single-branch", str(remote), str(repo)) + return remote, repo + + +def _resolve(repo: Path, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(ROOT / "scripts" / "default_branch.py"), *args], + cwd=repo, + capture_output=True, + text=True, + check=False, + ) + + +def _make(repo: Path, target: str, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["make", target, "LINT_DEP_INSTALL=", "LINT_DEP_BASE=", *args], + cwd=repo, + capture_output=True, + text=True, + check=False, + env={key: value for key, value in os.environ.items() if key != "BASE_REF"}, + ) + + +def test_existing_single_branch_clone_follows_remote_switch(remote_and_clone: tuple[Path, Path]) -> None: + remote, repo = remote_and_clone + before: Final = _resolve(repo) + assert before.returncode == 0, before.stderr + assert before.stdout.strip() == "origin/litellm_internal_staging" + _git(remote, "symbolic-ref", "HEAD", "refs/heads/main") + after: Final = _resolve(repo) + assert after.returncode == 0, after.stderr + assert after.stdout.strip() == "origin/main" + assert _git(repo, "rev-parse", "origin/main") == _git(remote, "rev-parse", "main") + assert _git(repo, "symbolic-ref", "refs/remotes/origin/HEAD").endswith("/litellm_internal_staging") + + +@pytest.mark.parametrize("missing_head", [False, True]) +def test_unverifiable_default_never_uses_cached_head( + remote_and_clone: tuple[Path, Path], + missing_head: bool, +) -> None: + remote, repo = remote_and_clone + if missing_head: + _git(remote, "symbolic-ref", "HEAD", "refs/heads/missing") + else: + _git(repo, "remote", "set-url", "origin", str(remote / "missing")) + result: Final = _resolve(repo) + assert result.returncode != 0 + assert not result.stdout + assert "explicit base ref" in result.stderr + checked: Final = _make(repo, "lint-format-check-changed") + assert checked.returncode != 0 + assert "No changed" not in checked.stdout + + +@pytest.mark.parametrize("base_ref", ["HEAD", "origin/litellm_internal_staging"]) +def test_explicit_base_works_without_remote_access( + remote_and_clone: tuple[Path, Path], + base_ref: str, +) -> None: + remote, repo = remote_and_clone + _git(repo, "remote", "set-url", "origin", str(remote / "missing")) + result: Final = _resolve(repo, "--base", base_ref) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == base_ref + checked: Final = _make(repo, "lint-format-check-changed", f"BASE_REF={base_ref}") + assert checked.returncode == 0, checked.stderr + assert "No changed litellm Python files" in checked.stdout + + +def test_budget_ratchet_compares_against_new_default(remote_and_clone: tuple[Path, Path]) -> None: + remote, repo = remote_and_clone + _git(remote, "symbolic-ref", "HEAD", "refs/heads/main") + resolved: Final = _resolve(repo) + assert resolved.returncode == 0, resolved.stderr + _git(repo, "checkout", "-qb", "litellm_feature", "origin/main") + (repo / "ruff-strict-budget.json").write_text('{"C901": {"limit": 1}}\n') + command: Final = [sys.executable, "scripts/budget_ratchet_check.py"] + checked: Final = subprocess.run(command, cwd=repo, capture_output=True, text=True, check=False) + assert checked.returncode == 1 + assert "limit raised 0 -> 1" in checked.stdout + assert "base origin/main" in checked.stdout + overridden: Final = subprocess.run( + [*command, "--base", "origin/litellm_internal_staging"], + cwd=repo, + capture_output=True, + text=True, + check=False, + ) + assert overridden.returncode == 0, overridden.stdout + overridden.stderr + + +def _freshness(repo: Path, *args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [ + sys.executable, + "-c", + "import sys; from pathlib import Path; " + "from ci_cd.run_migration import _check_branch_freshness; " + "_check_branch_freshness(Path(sys.argv[1]), sys.argv[2] if len(sys.argv) > 2 else None)", + str(repo), + *args, + ], + cwd=ROOT, + capture_output=True, + text=True, + check=False, + ) + + +def test_migration_freshness_refuses_stale_branch_after_switch(remote_and_clone: tuple[Path, Path]) -> None: + remote, repo = remote_and_clone + before: Final = _freshness(repo) + assert before.returncode == 0, before.stderr + assert "Branch freshness OK" in before.stdout + _git(remote, "symbolic-ref", "HEAD", "refs/heads/main") + after: Final = _freshness(repo) + assert after.returncode == 3 + assert "1 commit(s) behind origin/main" in after.stderr + overridden: Final = _freshness(repo, "litellm_internal_staging") + assert overridden.returncode == 0, overridden.stderr + _git(repo, "merge", "--ff-only", "origin/main") + updated: Final = _freshness(repo) + assert updated.returncode == 0, updated.stderr + assert "up to date with origin/main" in updated.stdout + + +def test_migration_freshness_refuses_unavailable_remote(remote_and_clone: tuple[Path, Path]) -> None: + remote, repo = remote_and_clone + _git(repo, "remote", "set-url", "origin", str(remote / "missing")) + result: Final = _freshness(repo) + assert result.returncode == 3 + assert "Could not discover origin's default branch" in result.stderr + explicit: Final = _freshness(repo, "litellm_internal_staging") + assert explicit.returncode == 3 + assert "git fetch origin litellm_internal_staging" in explicit.stderr + + +@pytest.mark.parametrize( + "gate", + [ + "budget_ratchet_check", + "ruff_strict_gate", + "type_discipline_gate", + "test_quality_gate", + "type_check_gate", + ], +) +def test_each_gate_refuses_an_unverifiable_default(remote_and_clone: tuple[Path, Path], gate: str) -> None: + remote, repo = remote_and_clone + _git(repo, "remote", "set-url", "origin", str(remote / "missing")) + result: Final = subprocess.run( + [sys.executable, f"scripts/{gate}.py"], + cwd=repo, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode != 0 + assert "Cannot verify the base branch against origin" in result.stderr + + +@pytest.mark.parametrize( + "target", ["lint-format-check-changed", "lint-test-quality", "lint-test-quality-budget-update"] +) +def test_direct_make_target_fetches_default_once(remote_and_clone: tuple[Path, Path], target: str) -> None: + _, repo = remote_and_clone + trace: Final = repo.parent / "git-trace.jsonl" + shutil.copyfile(ROOT / "scripts" / "check_test_quality.py", repo / "scripts" / "check_test_quality.py") + shutil.copyfile(ROOT / "test-quality-budget.json", repo / "test-quality-budget.json") + (repo / "tests").mkdir() + result: Final = subprocess.run( + ["make", "-o", "install-dev", target, "LINT_DEP_INSTALL=", "UV_RUN=env"], + cwd=repo, + capture_output=True, + text=True, + check=False, + env={**{key: value for key, value in os.environ.items() if key != "BASE_REF"}, "GIT_TRACE2_EVENT": str(trace)}, + ) + assert result.returncode == 0, result.stdout + result.stderr + commands: Final = tuple( + event["argv"][1:] + for line in trace.read_text().splitlines() + if (event := json.loads(line)).get("event") == "start" + ) + assert sum(command[0] == "ls-remote" for command in commands) == 1 + assert sum(command[0] == "fetch" for command in commands) == 1 diff --git a/tests/test_litellm/test_exception_header_preservation.py b/tests/test_litellm/test_exception_header_preservation.py index 6ea478c633b..dd142d9d40b 100644 --- a/tests/test_litellm/test_exception_header_preservation.py +++ b/tests/test_litellm/test_exception_header_preservation.py @@ -18,6 +18,7 @@ from litellm.exceptions import ( ImageFetchError, MidStreamFallbackError, RateLimitError, + ServiceUnavailableError, ) @@ -312,3 +313,85 @@ class TestProxyHeaderExtraction: # Verify headers are extracted and prefixed correctly assert headers.get("llm_provider-x-request-id") == "req-abc123" assert headers.get("llm_provider-x-ms-region") == "eastus" + + +class TestBedrockErrorHeaders: + """A BedrockError built with headers but no response still exposes them (LIT-5428).""" + + def test_synthesized_response_carries_headers(self): + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError( + status_code=500, + message="Amazon Bedrock is unable to process your request.", + headers={"x-amzn-RequestId": "req-base-500"}, + ) + + assert error.response.headers["x-amzn-requestid"] == "req-base-500" + assert str(error.request.url) == str(BedrockError(status_code=500, message="boom").request.url) + assert str(error.response.request.url) == str(error.request.url) + + def test_synthesized_response_without_headers_stays_empty(self): + from litellm.llms.bedrock.common_utils import BedrockError + + error = BedrockError(status_code=500, message="boom") + + assert dict(error.response.headers) == {} + + def test_explicit_response_is_kept(self): + from litellm.llms.bedrock.common_utils import BedrockError + + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "from-response"}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + error = BedrockError( + status_code=500, + message="boom", + headers={"x-amzn-RequestId": "from-headers"}, + response=provider_response, + ) + + assert error.response is provider_response + + def test_proxy_extraction_surfaces_bedrock_request_id(self): + """End-to-end shape the proxy error handler returns to the caller.""" + from litellm.litellm_core_utils.exception_mapping_utils import exception_type + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers, + ) + from litellm.llms.bedrock.common_utils import BedrockError + + provider_response = httpx.Response( + status_code=500, + headers={"x-amzn-RequestId": "req-proxy-500"}, + text='{"message":"Amazon Bedrock is unable to process your request."}', + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/"), + ) + + with pytest.raises(ServiceUnavailableError) as exc_info: + exception_type( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + original_exception=BedrockError( + status_code=500, + message=provider_response.text, + headers=provider_response.headers, + response=provider_response, + ), + custom_llm_provider="bedrock", + completion_kwargs={}, + extra_kwargs={}, + ) + + # Mirrors ProxyBaseLLMRequestProcessing._handle_llm_api_exception + error = exc_info.value + headers = getattr(error, "headers", None) or {} + if not headers: + _response = getattr(error, "response", None) + if _response is not None: + _response_headers = getattr(_response, "headers", None) + if _response_headers: + headers = get_response_headers(dict(_response_headers)) + + assert headers.get("llm_provider-x-amzn-requestid") == "req-proxy-500" diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index a7a9e0fc37d..701938f5677 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -14,26 +14,9 @@ import os import pytest -import litellm from litellm.utils import get_model_info -@pytest.fixture(scope="module", autouse=True) -def _local_model_cost_map(): - """ - Point litellm at the bundled cost map for the duration of this module - only. ``mp.undo()`` restores both the environment variable and - ``litellm.model_cost`` so nothing leaks into later tests. - """ - mp = pytest.MonkeyPatch() - mp.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - mp.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - get_model_info.cache_clear() - yield - mp.undo() - get_model_info.cache_clear() - - NEW_ENTRIES = { "fireworks_ai/accounts/fireworks/models/deepseek-v4-pro-0813": { "input_cost_per_token": 1.32e-06, @@ -54,19 +37,6 @@ def model_data(): return json.load(f) -def test_fireworks_serverless_entries_exist(model_data): - """The new prefixed entry carries the pricing and metadata from #37274.""" - for key, expected in NEW_ENTRIES.items(): - assert key in model_data, f"{key} is missing from model_prices_and_context_window.json" - entry = model_data[key] - for field, value in expected.items(): - assert entry[field] == pytest.approx(value), f"{key}.{field}" - assert entry["litellm_provider"] == "fireworks_ai" - assert entry["mode"] == "chat" - assert entry["supports_function_calling"] is True - assert entry["supports_vision"] is False - - def test_bare_fireworks_ids_resolve_through_prefixed_entries(): """Bare IDs from #37274 resolve via the provider-prefix lookup path.""" for bare_id, prefixed_key in [ diff --git a/tests/test_litellm/test_gate_slot_lock.py b/tests/test_litellm/test_gate_slot_lock.py index 17fa8547ce7..c80e876700b 100644 --- a/tests/test_litellm/test_gate_slot_lock.py +++ b/tests/test_litellm/test_gate_slot_lock.py @@ -1,6 +1,8 @@ import fcntl import importlib.util +import json import os +import shlex import signal import subprocess import sys @@ -301,33 +303,47 @@ def test_held_slot_context_manager_releases_on_exit(tmp_path: Path, monkeypatch: fcntl.flock(probe, fcntl.LOCK_UN) -def _make_rule(target: str) -> tuple[list[str], list[str]]: - database = subprocess.run( - ["make", "--dry-run", "--print-data-base", "info"], - cwd=ROOT, - capture_output=True, - text=True, - check=True, - ).stdout - lines = database.splitlines() - for index, line in enumerate(lines): - if line != f"{target}:" and not line.startswith(f"{target}: "): - continue - recipe: list[str] = [] - for follower in lines[index + 1 :]: - if follower.startswith("#"): - continue - if not follower.startswith("\t"): - break - recipe.append(follower.strip()) - return line.split(":", 1)[1].split(), recipe - raise AssertionError(f"target {target} not found in make database") - - -def test_direct_make_lint_takes_a_slot_before_any_setup() -> None: - lint_prerequisites, lint_recipe = _make_rule("lint") - assert lint_prerequisites == [] - assert any("$(GATE_SLOT_LOCK)" in line for line in lint_recipe) - inner_prerequisites, _ = _make_rule("lint-inner") - assert "lint-install" in inner_prerequisites - assert "lint-fetch-base" in inner_prerequisites +def test_direct_make_lint_takes_a_slot_before_any_setup(tmp_path: Path) -> None: + lock_dir = tmp_path / "locks" + lock_dir.mkdir() + events_file = tmp_path / "setup.jsonl" + stderr_file = tmp_path / "make.stderr" + probe = tmp_path / "probe.py" + probe.write_text( + "import fcntl, json, os, pathlib, sys\n" + "with (pathlib.Path(os.environ['LITELLM_GATE_SLOT_DIR']) / 'slot-0.lock').open('wb') as slot:\n" + " try:\n" + " fcntl.flock(slot, fcntl.LOCK_EX | fcntl.LOCK_NB)\n" + " locked = False\n" + " except BlockingIOError:\n" + " locked = True\n" + "with open(os.environ['EVENTS_FILE'], 'a') as events:\n" + " events.write(json.dumps({'phase': sys.argv[1], 'locked': locked}) + '\\n')\n" + "if sys.argv[1] == 'base':\n" + " print('HEAD')\n" + ) + (tmp_path / "Makefile").write_text((ROOT / "Makefile").read_text()) + command = [ + "make", "-o", "lint-checks", "lint", "MAKE=make -o lint-checks", + f"GATE_SLOT_LOCK={shlex.join([sys.executable, str(HELPER)])}", + f"UV={shlex.join([sys.executable, str(probe), 'setup'])}", + f"UV_RUN={shlex.join([sys.executable, str(probe), 'setup'])}", + f"RESOLVE_BASE={shlex.join([sys.executable, str(probe), 'base'])}", + ] + with (lock_dir / "slot-0.lock").open("wb") as held, stderr_file.open("wb") as stderr: + fcntl.flock(held, fcntl.LOCK_EX) + process = subprocess.Popen( + command, cwd=tmp_path, stdout=subprocess.DEVNULL, stderr=stderr, + env={**_env(lock_dir, "1"), "EVENTS_FILE": str(events_file)}, + ) + try: + assert _wait_until(lambda: "queueing" in stderr_file.read_text(), 10) + assert not events_file.exists() + fcntl.flock(held, fcntl.LOCK_UN) + assert process.wait(timeout=30) == 0, stderr_file.read_text() + finally: + fcntl.flock(held, fcntl.LOCK_UN) + _reap(process) + events = tuple(json.loads(line) for line in events_file.read_text().splitlines()) + assert {event["phase"] for event in events} == {"setup", "base"} + assert all(event["locked"] for event in events) diff --git a/tests/test_litellm/test_gpt_5_5_model_metadata.py b/tests/test_litellm/test_gpt_5_5_model_metadata.py index a60fa9466e6..e07efbcc913 100644 --- a/tests/test_litellm/test_gpt_5_5_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_5_model_metadata.py @@ -1,54 +1,6 @@ import json from pathlib import Path -import pytest - -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - -@pytest.mark.parametrize("model", ["azure_ai/gpt-5.5", "azure_ai/gpt-5.5-2026-04-23"]) -def test_azure_ai_gpt_5_5_model_info(model): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "azure_ai" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 5e-06 - assert info["output_cost_per_token"] == 3e-05 - assert info["cache_read_input_token_cost"] == 5e-07 - - assert info["input_cost_per_token_above_272k_tokens"] == 1e-05 - assert info["output_cost_per_token_above_272k_tokens"] == 4.5e-05 - assert info["cache_read_input_token_cost_above_272k_tokens"] == 1e-06 - - assert info["input_cost_per_token_priority"] == 1e-05 - assert info["output_cost_per_token_priority"] == 6e-05 - - assert info["max_input_tokens"] == 1050000 - assert info["max_output_tokens"] == 128000 - assert info["max_tokens"] == 128000 - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_web_search"] is True - # gpt-5.5 dropped minimal reasoning effort support (true on gpt-5.4) - assert info["supports_minimal_reasoning_effort"] is False - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.split("/", 1)[1] - assert provider == "azure_ai" - def test_azure_ai_gpt_5_5_backup_matches_main(): """Ensure the bundled model cost map stays in sync with the canonical file.""" diff --git a/tests/test_litellm/test_gpt_realtime_mode.py b/tests/test_litellm/test_gpt_realtime_mode.py index 314fd63c4cc..8c41e474486 100644 --- a/tests/test_litellm/test_gpt_realtime_mode.py +++ b/tests/test_litellm/test_gpt_realtime_mode.py @@ -1,10 +1,8 @@ import json from pathlib import Path -import pytest from typing_extensions import get_args, get_type_hints -import litellm from litellm.types.utils import ModelInfoBase REALTIME_ONLY_GPT_MODELS = ( @@ -43,44 +41,11 @@ REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS = ( ALL_REALTIME_ONLY_GPT_MODELS = REALTIME_ONLY_GPT_MODELS + REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS -def _load_cost_map() -> dict: - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - return json.load(f) - - def test_realtime_is_a_valid_mode_literal(): hints = get_type_hints(ModelInfoBase, include_extras=False) assert "realtime" in get_args(hints["mode"]) -@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS) -def test_realtime_only_gpt_models_are_mode_realtime(model): - """These models only serve /v1/realtime and are rejected by /v1/chat/completions - ("This is not a chat model ..."), so they must not be tagged mode=chat.""" - info = _load_cost_map()[model] - assert info["supported_endpoints"] == ["/v1/realtime"] - assert info["mode"] == "realtime" - - -@pytest.mark.parametrize("model", REALTIME_ONLY_GPT_MODELS_WITHOUT_ENDPOINTS) -def test_realtime_only_gpt_4o_models_are_mode_realtime(model): - """gpt-4o(-mini)-realtime-preview are realtime-only and must not be mode=chat.""" - assert _load_cost_map()[model]["mode"] == "realtime" - - -def test_get_model_info_reports_realtime_mode(monkeypatch): - """get_model_info must resolve the retag against the bundled cost map, not the - hosted map fetched from main, which lags this repo until the next promotion.""" - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - litellm.get_model_info.cache_clear() - try: - assert litellm.get_model_info("gpt-realtime-mini")["mode"] == "realtime" - finally: - litellm.get_model_info.cache_clear() - - def test_backup_matches_main_for_realtime_models(): repo_root = Path(__file__).parents[2] with open(repo_root / "model_prices_and_context_window.json") as f: diff --git a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py index 6f1ba702d8d..d73311baae9 100644 --- a/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py +++ b/tests/test_litellm/test_mistral_medium_3_5_model_metadata.py @@ -3,8 +3,6 @@ from pathlib import Path import pytest -import litellm -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider REPO_ROOT = Path(__file__).parents[2] MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -27,56 +25,6 @@ def _load(path): return json.load(f) - -@pytest.mark.parametrize("model", MEDIUM_3_5_MODELS) -def test_medium_3_5_specs(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "mistral" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - assert info["max_tokens"] == 262144 - - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_assistant_prefill"] is True - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.split("/", 1)[1] - assert provider == "mistral" - - -def test_mistral_medium_latest_resolves_to_medium_3_5(local_model_cost_map): - """LIT-3883: the -latest alias was retargeted to Medium 3.5; get_model_info must - return the 3.5 pricing/context/reasoning, not the stale Medium 3.1 values.""" - info = litellm.get_model_info(model="mistral/mistral-medium-latest") - - assert info["input_cost_per_token"] == 1.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - assert info["max_input_tokens"] == 262144 - assert info["supports_reasoning"] is True - - -def test_mistral_medium_2508_keeps_medium_3_1_specs(): - """The date-pinned 2508 alias is Medium 3.1 and must not inherit 3.5 pricing.""" - info = _load(MAIN_PATH).get("mistral/mistral-medium-2508") - assert info is not None, "mistral/mistral-medium-2508 missing from cost map" - - assert info["input_cost_per_token"] == 4e-07 - assert info["output_cost_per_token"] == 2e-06 - assert info["max_input_tokens"] == 131072 - assert info.get("supports_reasoning") is not True - - @pytest.mark.parametrize("model", SYNCED_MODELS) def test_backup_matches_main(model): """Ensure the bundled (backup) cost map stays in sync with the canonical file.""" diff --git a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py index 0442321ba0b..182c444bac9 100644 --- a/tests/test_litellm/test_mistral_small_4_0_model_metadata.py +++ b/tests/test_litellm/test_mistral_small_4_0_model_metadata.py @@ -18,29 +18,6 @@ def _load(path): return json.load(f) -@pytest.mark.parametrize("model", SMALL_4_0_MODELS) -def test_small_4_0_specs(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "mistral" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.5e-07 - assert info["output_cost_per_token"] == 6e-07 - - assert info["max_input_tokens"] == 262144 - assert info["max_output_tokens"] == 262144 - assert info["max_tokens"] == 262144 - - assert info["supports_reasoning"] is True - assert info["supports_vision"] is True - assert info["supports_function_calling"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_assistant_prefill"] is True - - @pytest.mark.parametrize("model", SMALL_4_0_MODELS) def test_backup_matches_main(model): main_cost = _load(MAIN_PATH) 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_muse_spark_1_2_model_metadata.py b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py index 0587883aa44..02527a98711 100644 --- a/tests/test_litellm/test_muse_spark_1_2_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_2_model_metadata.py @@ -23,46 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di return json.load(f) - -@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) -def test_muse_spark_1_2_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): - info = _load_cost_map().get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "meta" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cached_cost - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 131072 - assert info["max_tokens"] == 131072 - - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_pdf_input"] is True - assert info["supports_web_search"] is True - assert info["supports_minimal_reasoning_effort"] is True - assert info["supports_xhigh_reasoning_effort"] is True - - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] - assert info["supported_modalities"] == ["text", "image", "video"] - assert info["supported_output_modalities"] == ["text"] - - assert info["search_context_cost_per_query"] == { - "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, - } - - @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) def test_muse_spark_1_2_cost_per_token( local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float diff --git a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py index 1ecd9490f78..92b099fc780 100644 --- a/tests/test_litellm/test_muse_spark_1_3_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_3_model_metadata.py @@ -23,46 +23,6 @@ def _load_cost_map(filename: str = "model_prices_and_context_window.json") -> di return json.load(f) - -@pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) -def test_muse_spark_1_3_model_info(model: str, input_cost: float, cached_cost: float, output_cost: float): - info = _load_cost_map().get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "meta" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cached_cost - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 131072 - assert info["max_tokens"] == 131072 - - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_pdf_input"] is True - assert info["supports_web_search"] is True - assert info["supports_minimal_reasoning_effort"] is True - assert info["supports_xhigh_reasoning_effort"] is True - - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] - assert info["supported_modalities"] == ["text", "image", "video"] - assert info["supported_output_modalities"] == ["text"] - - assert info["search_context_cost_per_query"] == { - "search_context_size_high": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_low": WEB_SEARCH_COST_PER_QUERY, - "search_context_size_medium": WEB_SEARCH_COST_PER_QUERY, - } - - @pytest.mark.parametrize("model, input_cost, cached_cost, output_cost", PRICING) def test_muse_spark_1_3_cost_per_token( local_model_cost_map, model: str, input_cost: float, cached_cost: float, output_cost: float diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index b84cb8aa657..e12da0833dc 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -159,12 +159,12 @@ def _commit_all(repo: Path, message: str) -> None: ) -def _set_base_ref(repo: Path) -> None: - subprocess.run( - ["git", "update-ref", "refs/remotes/origin/litellm_internal_staging", "HEAD"], - cwd=repo, - check=True, - ) +def _set_base_ref(repo: Path, branch: str = "litellm_internal_staging") -> None: + remote = repo.parent / "remote.git" + subprocess.run(["git", "clone", "-q", "--bare", str(repo), str(remote)], check=True) + subprocess.run(["git", "update-ref", f"refs/heads/{branch}", "HEAD"], cwd=remote, check=True) + subprocess.run(["git", "symbolic-ref", "HEAD", f"refs/heads/{branch}"], cwd=remote, check=True) + subprocess.run(["git", "remote", "add", "origin", str(remote)], cwd=repo, check=True) def _stage_file(repo: Path, relative: str, body: str) -> None: @@ -174,10 +174,11 @@ def _stage_file(repo: Path, relative: str, body: str) -> None: 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: +@pytest.mark.parametrize("branch", ["litellm_internal_staging", "main"]) +def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path, branch: str) -> None: repo, bin_dir = _sandbox(tmp_path) _commit_all(repo, "base") - _set_base_ref(repo) + _set_base_ref(repo, branch) (repo / "litellm" / "foo.py").write_text("x = 2\n") proc = _run(repo, bin_dir, {}) assert proc.returncode == 0, proc.stdout + proc.stderr @@ -261,8 +262,8 @@ def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Pat _commit_all(repo, "base") proc = _run(repo, bin_dir, {}) assert proc.returncode == 1 - assert "cannot resolve the merge base" in proc.stdout - assert "git fetch origin litellm_internal_staging" in proc.stdout + assert "Cannot verify the base branch against origin" in proc.stdout + assert "explicit base ref" in proc.stdout assert "check: FAIL" in proc.stdout @@ -622,3 +623,28 @@ def test_failing_run_ends_with_a_fail_verdict(tmp_path: Path) -> None: assert proc.returncode == 1 assert "check: FAIL" in proc.stdout assert "check: PASS" not in proc.stdout + + + +def test_explicit_base_scopes_offline_without_a_remote(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + (repo / "litellm" / "foo.py").write_text("x = 2\n") + proc = _run(repo, bin_dir, {"BASE_REF": "HEAD"}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "merge base with HEAD" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_symlinked_hook_can_resolve_default_branch(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo, "main") + hook = repo / ".git" / "hooks" / "pre-commit" + hook.symlink_to(SCRIPT) + proc = subprocess.run( + [str(hook)], cwd=repo, capture_output=True, text=True, + env=_env(repo, bin_dir, {}), timeout=120, + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "no branch changes vs origin/main" in proc.stdout diff --git a/tests/test_litellm/test_replicate_model_key_format.py b/tests/test_litellm/test_replicate_model_key_format.py index 8c2b72f8ed2..77ae5e1b069 100644 --- a/tests/test_litellm/test_replicate_model_key_format.py +++ b/tests/test_litellm/test_replicate_model_key_format.py @@ -20,14 +20,6 @@ def test_replicate_models_have_valid_key_prefix(model_cost: dict[str, Any]) -> N ) -def test_replicate_openai_gpt_oss_20b_key_exists(model_cost: dict[str, Any]) -> None: - assert "replicate/openai/gpt-oss-20b" in model_cost - info = model_cost["replicate/openai/gpt-oss-20b"] - assert info["litellm_provider"] == "replicate" - assert info["mode"] == "chat" - assert info["supports_function_calling"] is True - - def test_replicate_backup_matches_main() -> None: repo_root = Path(__file__).parents[2] main_path = repo_root / "model_prices_and_context_window.json" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index e9b77076448..eed34c79a06 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -2387,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. diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py index c9e2863d240..b9764eca2f8 100644 --- a/tests/test_litellm/test_together_ai_model_metadata.py +++ b/tests/test_litellm/test_together_ai_model_metadata.py @@ -5,7 +5,6 @@ from typing import Final import pytest from pydantic import TypeAdapter -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider REPO_ROOT: Final = Path(__file__).parents[2] @@ -77,59 +76,6 @@ def cost_map() -> CostMap: return COST_MAP_ADAPTER.validate_python(json.load(f)) -@pytest.mark.parametrize("model", SERVERLESS_CHAT_MODELS) -def test_together_serverless_chat_model_is_mapped(cost_map: CostMap, model: str): - info = cost_map.get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - assert info["litellm_provider"] == "together_ai" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] >= 0 - assert info["output_cost_per_token"] >= info["input_cost_per_token"] - assert "deprecation_date" not in info - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.removeprefix("together_ai/") - assert provider == "together_ai" - - -def test_together_kimi_k3_pricing_and_capabilities(cost_map: CostMap): - info = cost_map["together_ai/moonshotai/Kimi-K3"] - assert info["input_cost_per_token"] == 3e-06 - assert info["output_cost_per_token"] == 1.5e-05 - assert info["max_input_tokens"] == 1048576 - assert info["supports_function_calling"] is True - assert info["supports_tool_choice"] is True - assert info["supports_response_schema"] is True - assert info["supports_vision"] is True - assert info["supports_reasoning"] is True - - -def test_together_glm_52_pricing(cost_map: CostMap): - info = cost_map["together_ai/zai-org/GLM-5.2"] - assert info["input_cost_per_token"] == 1.4e-06 - assert info["output_cost_per_token"] == 4.4e-06 - assert info["max_input_tokens"] == 1048575 - assert info["max_output_tokens"] == 128000 - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - - -def test_together_glm_53_flash_pricing_and_capabilities(cost_map: CostMap): - info = cost_map["together_ai/zai-org/GLM-5.3-Flash"] - assert info["input_cost_per_token"] == 1.5e-07 - assert info["output_cost_per_token"] == 5e-07 - assert info["cache_read_input_token_cost"] == 3e-08 - assert info["max_input_tokens"] == 1048575 - assert info["max_output_tokens"] == 128000 - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_tool_choice"] is True - assert info["supports_response_schema"] is True - assert info["supports_vision"] is True - assert info["supports_reasoning"] is True - - def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost_map: CostMap): inflated = sorted( model @@ -142,21 +88,6 @@ def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost assert inflated == [] -def test_together_multilingual_e5_embedding_entry(cost_map: CostMap): - info = cost_map["together_ai/intfloat/multilingual-e5-large-instruct"] - assert info["mode"] == "embedding" - assert info["input_cost_per_token"] == 2e-08 - assert info["max_input_tokens"] == 514 - assert info["output_vector_size"] == 1024 - - -def test_together_llama_33_70b_repriced_to_current_together_rate(cost_map: CostMap): - info = cost_map["together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo"] - assert info["input_cost_per_token"] == 1.04e-06 - assert info["output_cost_per_token"] == 1.04e-06 - assert info["max_input_tokens"] == 131072 - - @pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS)) def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str): info = cost_map.get(model) @@ -210,32 +141,7 @@ CACHED_INPUT_MODELS: Final = ( ) -@pytest.mark.parametrize("model", CACHED_INPUT_MODELS) -def test_together_cached_input_model_carries_cache_read_pricing(cost_map: CostMap, model: str): - info = cost_map.get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - assert info.get("supports_prompt_caching") is True - cache_read = info.get("cache_read_input_token_cost") - assert isinstance(cache_read, float) - assert 0 < cache_read < info["input_cost_per_token"] - assert "cache_creation_input_token_cost" not in info - - def test_together_prompt_caching_flag_implies_cache_read_rate(cost_map: CostMap): for model, info in cost_map.items(): if model.startswith("together_ai/") and info.get("supports_prompt_caching"): assert "cache_read_input_token_cost" in info, f"{model} flags caching without a cache read rate" - - -def test_together_deepseek_v4_flash_cache_read_rate(cost_map: CostMap): - info = cost_map["together_ai/deepseek-ai/DeepSeek-V4-Flash-0731"] - assert info["input_cost_per_token"] == 1.4e-07 - assert info["cache_read_input_token_cost"] == 3e-08 - assert info["output_cost_per_token"] == 2.8e-07 - - -def test_together_qwen_37_max_repriced_to_current_together_rate(cost_map: CostMap): - info = cost_map["together_ai/Qwen/Qwen3.7-Max"] - assert info["input_cost_per_token"] == 2.5e-06 - assert info["output_cost_per_token"] == 7.5e-06 - assert info["cache_read_input_token_cost"] == 5e-07 diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 271c84384b4..4ce68b71079 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` @@ -6087,3 +6092,47 @@ class TestFinalOptionalParamsLineRedaction: assert "'max_tokens': 17" in printed assert "'temperature': 0.25" in printed + + +def _credential_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [record.getMessage() for record in caplog.records if "litellm_credential_name=" in record.getMessage()] + + +def test_load_credentials_from_list_warns_when_the_named_credential_is_not_loaded( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.utils import load_credentials_from_list + + monkeypatch.setattr(litellm, "credential_list", []) + request_kwargs = {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"} + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + load_credentials_from_list(request_kwargs) + + assert request_kwargs == {"litellm_credential_name": "openai-cred", "model": "openai/gpt-5.4-mini"} + assert _credential_warnings(caplog) == [ + "litellm_credential_name=openai-cred matched none of the 0 loaded credentials; the request runs without it" + ] + + +def test_load_credentials_from_list_fills_kwargs_from_the_loaded_credential_without_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.types.utils import CredentialItem + from litellm.utils import load_credentials_from_list + + loaded = CredentialItem( + credential_name="openai-cred", + credential_values={"api_key": "sk-from-db", "api_base": "https://credential.example"}, + credential_info={}, + ) + monkeypatch.setattr(litellm, "credential_list", [loaded]) + request_kwargs = {"litellm_credential_name": "openai-cred", "api_base": "https://request.example"} + with caplog.at_level(logging.WARNING, logger=verbose_logger.name): + load_credentials_from_list(request_kwargs) + + assert request_kwargs == { + "litellm_credential_name": "openai-cred", + "api_base": "https://request.example", + "api_key": "sk-from-db", + } + assert _credential_warnings(caplog) == [] 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_rust/conftest.py b/tests/test_litellm_rust/conftest.py new file mode 100644 index 00000000000..02d274bb405 --- /dev/null +++ b/tests/test_litellm_rust/conftest.py @@ -0,0 +1,24 @@ +import os + +import pytest + + +def pytest_collection_modifyitems(items): + rust_enabled = os.environ.get("LITELLM_RUST", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + if not rust_enabled: + skip = pytest.mark.skip(reason="requires LITELLM_RUST=1 and a compiled Rust extension") + for item in items: + item.add_marker(skip) + return + + try: + from litellm.rust_bridge import _native # noqa: F401 # validates the installed extension + except ImportError as error: + raise pytest.UsageError( + "LITELLM_RUST=1 requires a compiled litellm.rust_bridge._native extension" + ) from error diff --git a/tests/test_litellm_rust/test_ocr.py b/tests/test_litellm_rust/test_ocr.py new file mode 100644 index 00000000000..d5b1fce1139 --- /dev/null +++ b/tests/test_litellm_rust/test_ocr.py @@ -0,0 +1,72 @@ +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest + +import litellm + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.fixture +def ocr_server(): + requests = [] + + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + requests.append( + { + "headers": {name.lower(): value for name, value in self.headers.items()}, + "body": json.loads(self.rfile.read(int(self.headers["Content-Length"]))), + } + ) + if self.headers.get("User-Agent", "").startswith("python-httpx"): + self.send_response(418) + self.end_headers() + return + response = json.dumps( + { + "pages": [{"index": 0, "markdown": "native OCR response", "images": [], "dimensions": None}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, + } + ).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response))) + self.end_headers() + self.wfile.write(response) + + def log_message(self, format, *args): + pass + + server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=lambda: server.serve_forever(poll_interval=0.01), daemon=True) + thread.start() + try: + yield server, requests + finally: + server.shutdown() + server.server_close() + thread.join() + + +def test_ocr_with_rust_extension(ocr_server): + server, requests = ocr_server + host, port = server.server_address + + response = litellm.ocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + api_key="test-key", + api_base=f"http://{host}:{port}", + ) + + assert response.pages[0].markdown == "native OCR response" + assert len(requests) == 1 + assert not requests[0]["headers"].get("user-agent", "").startswith("python-httpx") + assert requests[0]["body"] == { + "model": "mistral-ocr-latest", + "document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}, + } diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e7186dfe186..0c0952289e2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22180 + "limit": 22174 }, "LIT002": { - "limit": 26729 + "limit": 26715 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16426 + "limit": 16398 }, "LIT011": { - "limit": 5506 + "limit": 5504 }, "LIT012": { "limit": 4486 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts index 442dcd48f66..e0a5271366f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_config.ts @@ -313,6 +313,26 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { return agentData; }; +export const parseMcpPermissionsForForm = (agent: any) => ({ + allowed_mcp_servers_and_groups: { + servers: agent.object_permission?.mcp_servers ?? [], + accessGroups: agent.object_permission?.mcp_access_groups ?? [], + toolsets: agent.object_permission?.mcp_toolsets ?? [], + }, + mcp_tool_permissions: agent.object_permission?.mcp_tool_permissions ?? {}, +}); + +/** + * Always includes every MCP key (empty when cleared) so removals persist; + * the proxy merges object_permission per key, leaving non-MCP grants untouched. + */ +export const buildMcpObjectPermission = (values: any) => ({ + mcp_servers: values.allowed_mcp_servers_and_groups?.servers ?? [], + mcp_access_groups: values.allowed_mcp_servers_and_groups?.accessGroups ?? [], + mcp_toolsets: values.allowed_mcp_servers_and_groups?.toolsets ?? [], + mcp_tool_permissions: values.mcp_tool_permissions ?? {}, +}); + /** * Parse agent data for form fields */ @@ -356,5 +376,6 @@ export const parseAgentForForm = (agent: any) => { : [], // extra_headers: already an array of strings extra_headers: agent.extra_headers ?? [], + ...parseMcpPermissionsForForm(agent), }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index de1eb153f6f..357f924cbe7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -1,4 +1,5 @@ import React from "react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; @@ -10,6 +11,12 @@ vi.mock("@/components/networking", () => ({ getAgentInfo: vi.fn(), patchAgentCall: vi.fn(), getAgentCreateMetadata: vi.fn(), + getProxyBaseUrl: vi.fn(() => ""), + getUiConfig: vi.fn(async () => ({})), + fetchMCPServers: vi.fn(async () => []), + fetchMCPAccessGroups: vi.fn(async () => []), + fetchMCPToolsets: vi.fn(async () => []), + listMCPTools: vi.fn(async () => ({ tools: [] })), })); vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ @@ -111,7 +118,14 @@ const bedrockAgentcoreInfo: AgentCreateInfo = { const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); -const renderView = () => render(); +const renderView = () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render( + + + , + ); +}; const openEditor = async (user: ReturnType) => { await user.click(await screen.findByRole("tab", { name: "Settings" })); @@ -161,6 +175,7 @@ describe("AgentInfoView update payload", () => { rpm_limit: 222, session_tpm_limit: 333, session_rpm_limit: 444, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); @@ -201,6 +216,7 @@ describe("AgentInfoView update payload", () => { rpm_limit: 222, session_tpm_limit: 333, session_rpm_limit: 444, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); @@ -278,9 +294,30 @@ describe("AgentInfoView update payload", () => { api_base: "https://other.example.com", model: "langgraph/asst_1", }, + object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, }); }); + it("keeps the agent's existing MCP grants in the update payload", async () => { + const existingMcpGrants = { + mcp_servers: ["srv-1"], + mcp_access_groups: ["grp-a"], + mcp_toolsets: ["toolset-1"], + mcp_tool_permissions: { "srv-1": ["tool_x"] }, + }; + vi.mocked(networking.getAgentInfo).mockResolvedValue({ + ...A2A_AGENT, + object_permission: existingMcpGrants, + } as never); + const user = setup(); + renderView(); + await openEditor(user); + + await save(user); + + expect(patchedPayload().object_permission).toEqual(existingMcpGrants); + }); + it("preserves the full AgentCore runtime ARN (including the resource id after runtime/) across an unedited save", async () => { vi.mocked(networking.getAgentCreateMetadata).mockResolvedValue([bedrockAgentcoreInfo]); vi.mocked(networking.getAgentInfo).mockResolvedValue(BEDROCK_AGENTCORE_AGENT as never); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx index 0936b8e13db..29d97a20afe 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx @@ -24,6 +24,18 @@ vi.mock("./agent_form_fields", () => ({ unmountedA2AFieldNames: () => [], })); +vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ + useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "github" }] }), +})); + +vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ + default: () =>
, +})); + +vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ + default: () =>
, +})); + const agent = { agent_id: "agent-1", agent_name: "support-agent", @@ -62,5 +74,18 @@ describe("AgentInfoView settings", () => { expect(token).toBe("sk-test"); expect(agentId).toBe("agent-1"); expect(payload.tpm_limit).toBe(42); + const clearedMcpGrants = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }; + expect(payload.object_permission).toEqual(clearedMcpGrants); + }); + + it("shows MCP grants with server names on the overview tab", async () => { + vi.mocked(networking.getAgentInfo).mockResolvedValue({ + ...agent, + object_permission: { mcp_servers: ["srv-1"] }, + } as unknown as Agent); + + render(); + + expect(await screen.findByText("github (srv-1)")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index eddeeec674b..6cb99e9692f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -15,16 +15,27 @@ import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } import { Agent } from "@/components/agents/types"; import { KeyResponse } from "@/components/key_team_helpers/key_list"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import KeyInfoView from "@/components/templates/key_info_view"; +import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; +import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; import AgentVirtualKeys from "./agent_virtual_keys"; import AgentFormFields, { unmountedA2AFieldNames } from "./agent_form_fields"; import DynamicAgentFormFields, { buildDynamicAgentData, unmountedDynamicFieldNames } from "./dynamic_agent_form_fields"; -import { AGENT_FORM_CONFIG, buildAgentDataFromForm, parseAgentForForm } from "./agent_config"; +import { + AGENT_FORM_CONFIG, + buildAgentDataFromForm, + buildMcpObjectPermission, + parseAgentForForm, + parseMcpPermissionsForForm, +} from "./agent_config"; import { AgentFormField, AgentFormValues, AgentNumberInput, AgentRequestPayload, + McpServerSelection, + labelWithHint, omitFieldValues, useCollapsiblePanels, } from "./AgentFormKit"; @@ -111,7 +122,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT } else { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset(parseDynamicAgentForForm(data, typeInfo)); + form.reset({ ...parseDynamicAgentForForm(data, typeInfo), ...parseMcpPermissionsForForm(data) }); } else { form.reset(parseAgentForForm(data)); } @@ -131,7 +142,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT if (agentType !== "a2a") { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset(parseDynamicAgentForForm(agent, typeInfo)); + form.reset({ ...parseDynamicAgentForForm(agent, typeInfo), ...parseMcpPermissionsForForm(agent) }); } } } @@ -139,6 +150,14 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT const selectedAgentTypeInfo = agentTypeMetadata.find((t) => t.agent_type === detectedAgentType); const watchedFormValues = useWatch({ control: form.control }); + const mcpSelection = useWatch({ control: form.control, name: "allowed_mcp_servers_and_groups" }); + const mcpToolPermissions = useWatch({ control: form.control, name: "mcp_tool_permissions" }); + const { data: mcpServers = [] } = useMCPServers(); + + const mcpServerLabel = (serverId: string) => { + const server = mcpServers.find((s) => s.server_id === serverId); + return server?.server_name ? `${server.server_name} (${serverId})` : serverId; + }; const discoveryRequest = useMemo( () => buildDiscoveryRequest(detectedAgentType, watchedFormValues || {}, selectedAgentTypeInfo), @@ -199,7 +218,10 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT ? overlayDiscoveredCardParams(built, appliedDiscoveredSelection.selected_card) : built; - await patchAgentCall(accessToken, agentId, updateData); + await patchAgentCall(accessToken, agentId, { + ...updateData, + object_permission: buildMcpObjectPermission(values), + }); toast.success("Agent updated successfully"); setIsEditing(false); fetchAgentInfo(); @@ -337,13 +359,20 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.object_permission && (agent.object_permission.mcp_servers?.length || agent.object_permission.mcp_access_groups?.length || + agent.object_permission.mcp_toolsets?.length || (agent.object_permission.mcp_tool_permissions && Object.keys(agent.object_permission.mcp_tool_permissions).length > 0)) && (

MCP Tool Permissions

{agent.object_permission.mcp_servers && agent.object_permission.mcp_servers.length > 0 && ( - {agent.object_permission.mcp_servers.join(", ")} + +
+ {agent.object_permission.mcp_servers.map((serverId) => ( +
{mcpServerLabel(serverId)}
+ ))} +
+ )} {agent.object_permission.mcp_access_groups && agent.object_permission.mcp_access_groups.length > 0 && ( @@ -351,13 +380,16 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.object_permission.mcp_access_groups.join(", ")} )} + {agent.object_permission.mcp_toolsets && agent.object_permission.mcp_toolsets.length > 0 && ( + {agent.object_permission.mcp_toolsets.join(", ")} + )} {agent.object_permission.mcp_tool_permissions && Object.keys(agent.object_permission.mcp_tool_permissions).length > 0 && (
{Object.entries(agent.object_permission.mcp_tool_permissions).map(([serverId, tools]) => (
- {serverId}:{" "} + {mcpServerLabel(serverId)}:{" "} {Array.isArray(tools) ? tools.join(", ") : String(tools)}
))} @@ -457,6 +489,41 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {rateLimitField("session_rpm_limit", "Session RPM Limit")}
+ +

MCP Servers

+ + + {({ value, onChange }) => ( + + )} + + +
+ ) => + form.setValue("mcp_tool_permissions", toolPerms) + } + /> +
+
+ {canCreateVectorStores && ( + + )}
; + const { accessToken, userRole, userId, isViewOnly } = useAuthorized(); + return ( + + ); } diff --git a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx index 30ce62d8081..a663e16c2b2 100644 --- a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx @@ -1,43 +1,19 @@ "use client"; -import { Suspense, useEffect } from "react"; -import { useRouter, useSearchParams } from "next/navigation"; +import { Suspense } from "react"; import { useChatShell } from "@/contexts/ChatShellContext"; -import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; -import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; +import ConnectFlowSurface from "@/components/chat/ConnectFlowSurface"; // useSearchParams() requires a Suspense boundary for static export. function IntegrationsPageContent() { const { accessToken, selectedMCPServers, setSelectedMCPServers } = useChatShell(); - const router = useRouter(); - const searchParams = useSearchParams(); - const oauthReturn = searchParams.get("mcpOauthReturn"); - // Set by the gateway DCR authorize when a DCR client sends the user here to - // authorize servers before finishing sign-in (see gateway_dcr_flow.py). The - // handle keys the sealed per-flow cookie; connect_client is the client origin - // for display only. connect_flow is NOT cleaned from the URL: the finish form - // needs it, and the sealed cookie (not the URL) is the security boundary. - const connectFlow = searchParams.get("connect_flow"); - const connectClient = searchParams.get("connect_client"); - - // Clean up the OAuth return param after it's been consumed — real routing means - // we no longer need it to pick a tab, but it should not linger in the address bar. - useEffect(() => { - if (oauthReturn) { - const url = new URL(window.location.href); - url.searchParams.delete("mcpOauthReturn"); - router.replace(url.pathname + url.search); - } - }, [oauthReturn, router]); return (
- {connectFlow && } -
); diff --git a/ui/litellm-dashboard/src/app/connect/page.test.tsx b/ui/litellm-dashboard/src/app/connect/page.test.tsx index 7d49a8b6a4c..6a6ba24bd87 100644 --- a/ui/litellm-dashboard/src/app/connect/page.test.tsx +++ b/ui/litellm-dashboard/src/app/connect/page.test.tsx @@ -2,101 +2,29 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { render, screen } from "@testing-library/react"; import ConnectPage from "./page"; -interface PanelProps { +interface SurfaceProps { accessToken: string; selectedServers: string[]; onChange: (servers: string[]) => void; - connectMode?: boolean; } -interface BannerProps { - flowHandle: string; - clientOrigin: string | null; -} - -const { mockReplace, mockPanel, mockBanner, state } = vi.hoisted(() => { - const state = { - oauthReturn: null as string | null, - connectFlow: null as string | null, - connectClient: null as string | null, - }; - return { - state, - mockReplace: vi.fn(), - mockPanel: vi.fn((_props: PanelProps) =>
), - mockBanner: vi.fn((_props: BannerProps) =>
), - }; -}); - -vi.mock("next/navigation", () => ({ - useRouter: () => ({ replace: mockReplace }), - useSearchParams: () => ({ - get: (key: string) => { - if (key === "mcpOauthReturn") return state.oauthReturn; - if (key === "connect_flow") return state.connectFlow; - if (key === "connect_client") return state.connectClient; - return null; - }, - }), +const { mockSurface } = vi.hoisted(() => ({ + mockSurface: vi.fn((_props: SurfaceProps) =>
), })); + vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "token-123" }), })); -vi.mock("@/components/chat/MCPAppsPanel", () => ({ default: mockPanel })); -vi.mock("@/components/chat/ConnectFlowBanner", () => ({ default: mockBanner })); +vi.mock("@/components/chat/ConnectFlowSurface", () => ({ default: mockSurface })); describe("ConnectPage", () => { afterEach(() => { - state.oauthReturn = null; - state.connectFlow = null; - state.connectClient = null; - mockReplace.mockClear(); - mockPanel.mockClear(); - mockBanner.mockClear(); + mockSurface.mockClear(); }); - it("renders the MCP connect panel with the user's access token", () => { + it("renders the gateway connect surface with the user's access token and an empty selection", () => { render(); - expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument(); - expect(mockPanel.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] }); - }); - - it("strips the mcpOauthReturn param from the URL after an OAuth return", () => { - state.oauthReturn = "apps"; - window.history.replaceState({}, "", "/connect?mcpOauthReturn=apps"); - render(); - expect(mockReplace).toHaveBeenCalledWith("/connect"); - }); - - it("does not rewrite the URL when there is no OAuth return param", () => { - render(); - expect(mockReplace).not.toHaveBeenCalled(); - }); - - it("mounts the gateway connect banner and puts the panel in connect mode for a DCR flow", () => { - state.connectFlow = "flow-handle-123"; - state.connectClient = "https://claude.ai"; - render(); - expect(screen.getByTestId("connect-flow-banner")).toBeInTheDocument(); - expect(mockBanner.mock.calls[0][0]).toMatchObject({ - flowHandle: "flow-handle-123", - clientOrigin: "https://claude.ai", - }); - expect(mockPanel.mock.calls[0][0].connectMode).toBe(true); - }); - - it("shows no connect banner and leaves connect mode off for a plain visit", () => { - render(); - expect(screen.queryByTestId("connect-flow-banner")).not.toBeInTheDocument(); - expect(mockBanner).not.toHaveBeenCalled(); - expect(mockPanel.mock.calls[0][0].connectMode).toBe(false); - }); - - it("keeps the connect flow handle in the URL while stripping the OAuth return param", () => { - state.oauthReturn = "apps"; - state.connectFlow = "flow-handle-123"; - window.history.replaceState({}, "", "/connect?connect_flow=flow-handle-123&mcpOauthReturn=apps"); - render(); - expect(mockReplace).toHaveBeenCalledWith("/connect?connect_flow=flow-handle-123"); + expect(screen.getByTestId("connect-flow-surface")).toBeInTheDocument(); + expect(mockSurface.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] }); }); }); diff --git a/ui/litellm-dashboard/src/app/connect/page.tsx b/ui/litellm-dashboard/src/app/connect/page.tsx index 3f0c269e86b..652c044197b 100644 --- a/ui/litellm-dashboard/src/app/connect/page.tsx +++ b/ui/litellm-dashboard/src/app/connect/page.tsx @@ -1,36 +1,19 @@ "use client"; -import { Suspense, useEffect, useState } from "react"; -import { useRouter, useSearchParams } from "next/navigation"; +import { Suspense, useState } from "react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; -import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; +import ConnectFlowSurface from "@/components/chat/ConnectFlowSurface"; function ConnectPageContent() { const { accessToken } = useAuthorized(); const [selectedServers, setSelectedServers] = useState([]); - const router = useRouter(); - const searchParams = useSearchParams(); - const oauthReturn = searchParams.get("mcpOauthReturn"); - const connectFlow = searchParams.get("connect_flow"); - const connectClient = searchParams.get("connect_client"); - - useEffect(() => { - if (oauthReturn) { - const url = new URL(window.location.href); - url.searchParams.delete("mcpOauthReturn"); - router.replace(url.pathname + url.search); - } - }, [oauthReturn, router]); return (
- {connectFlow && } -
); diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 15cfff01766..111eeebe5a3 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -44,8 +44,9 @@ import { } from "./ComplexityRouterConfig"; const DEFAULT_SCORING_EXPLANATION = - "The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical " + - "terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the tier:"; + "The router scores each request across 7 built-in dimensions: token count, code presence, reasoning markers, technical " + + "terms, simple indicators, multi-step patterns, and question complexity, plus any custom dimensions you add. " + + "The weighted score determines the tier:"; const HEURISTIC_V2_EXPLANATION = "The router estimates success probability for all four tiers with the bundled calibrated model, then selects " + diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 2970e14b335..0c12cc0ba1a 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -1521,14 +1521,16 @@ describe("ComplexityRouterConfig tier editing", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.queryByText("How Classification Works")).not.toBeInTheDocument(); - expect(screen.queryByText("scores each request across 7 dimensions", { exact: false })).not.toBeInTheDocument(); + expect( + screen.queryByText("scores each request across 7 built-in dimensions", { exact: false }), + ).not.toBeInTheDocument(); }); it("keeps the scorer card on a built-in router, whose tiers the score still decides", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("How Classification Works")).toBeInTheDocument(); - expect(screen.getByText("scores each request across 7 dimensions", { exact: false })).toBeInTheDocument(); + expect(screen.getByText("scores each request across 7 built-in dimensions", { exact: false })).toBeInTheDocument(); }); it("says why a custom row is blocked instead of only reddening its border", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index d7df6ce33bb..de80e714e66 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -50,6 +50,7 @@ import EscalationKeywords from "./EscalationKeywords"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs"; +import { type CustomDimensionRow } from "./custom_dimensions"; import CompressionControls from "./CompressionControls"; import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression"; @@ -432,6 +433,11 @@ export interface ComplexityRouterConfigValue { tier_boundaries?: TierBoundaries; token_thresholds?: TokenThresholds; dimension_weights?: DimensionWeights; + /** + * Operator-added scoring dimensions, each carrying its own inline weight. Undefined means the router has + * none and keeps the key out of the payload; an empty array is a real "the last row was removed" state. + */ + custom_dimensions?: CustomDimensionRow[]; /** * Score floor the reasoning-marker override must clear. Undefined keeps the key out of the payload, so the * floor tracks tier_boundaries.simple_medium; an explicit 0 is a real floor that promotes on the markers alone. diff --git a/ui/litellm-dashboard/src/components/add_model/CustomDimensionRows.tsx b/ui/litellm-dashboard/src/components/add_model/CustomDimensionRows.tsx new file mode 100644 index 00000000000..34d8370aafb --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/CustomDimensionRows.tsx @@ -0,0 +1,136 @@ +import { Trash2 } from "lucide-react"; +import { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Textarea } from "@/components/ui/textarea"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Slider } from "@/components/ui/slider"; +import type { CustomDimensionRow } from "./custom_dimensions"; + +const SCORING_MODES = [ + { value: "binary", label: "Binary" }, + { value: "match_count", label: "Match count" }, +] as const; + +interface Props { + rows: CustomDimensionRow[]; + disabled: boolean; + onChange: (rows: CustomDimensionRow[]) => void; + onWeight: (id: string, weight: number) => void; + onAdd: () => void; + onRemove: (id: string) => void; +} + +export default function CustomDimensionRows({ rows, disabled, onChange, onWeight, onAdd, onRemove }: Props) { + const [draft, setDraft] = useState<{ id: string; raw: string } | null>(null); + const update = (id: string, patch: Partial) => + onChange(rows.map((row) => (row.id === id ? { ...row, ...patch } : row))); + const editWeight = (id: string, raw: string) => { + setDraft({ id, raw }); + if (raw.trim() && Number.isFinite(Number(raw))) onWeight(id, Number(raw)); + }; + return ( +
+ {rows.map((row, index) => ( +
+ Custom dimension {index + 1} +
+ +
+
+ + update(row.id, { name: event.target.value })} + /> +
+
+ + onWeight(row.id, Array.isArray(value) ? value[0] : value)} + /> + setDraft(null)} + onChange={(event) => editWeight(row.id, event.target.value)} + /> +
+
+ {(["keywords", "patterns"] as const).map((field) => ( +
+ +