diff --git a/.github/actions/cache-prisma-binaries/action.yml b/.github/actions/cache-prisma-binaries/action.yml new file mode 100644 index 00000000000..68615e94c08 --- /dev/null +++ b/.github/actions/cache-prisma-binaries/action.yml @@ -0,0 +1,40 @@ +name: "Cache Prisma binaries" +description: >- + Cache the Prisma CLI and engine binaries that `prisma generate` downloads, so + only the first job on a given prisma-client-py version pays for the download. + + prisma-client-py shells out to `npm install prisma@` whenever its + binary cache directory has no CLI entrypoint, which pulls ~85 MB of query and + schema engines over the network. That normally takes a few seconds, but it is + unbounded: one shard of a proxy-db run took 5m18s on that single step versus + 3.8s on its eleven siblings, which pushed the job past its timeout and got a + fully passing test run cancelled. + + Callers must not set PRISMA_BINARY_CACHE_DIR. The prisma-client-py default + (~/.cache/prisma-python/binaries//) is already + keyed by both versions, so a cache entry can never be served to a run that + expects different binaries. + +runs: + using: composite + steps: + - name: Resolve prisma-client-py version + id: version + shell: bash + run: | + version="$(grep -A1 '^name = "prisma"$' uv.lock | sed -n 's/^version = "\(.*\)"$/\1/p' | head -1)" + if [ -z "${version}" ]; then + echo "could not resolve the prisma package version from uv.lock" >&2 + exit 1 + fi + echo "version=${version}" >> "$GITHUB_OUTPUT" + + - name: Restore Prisma binaries + uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 + with: + # ~/.cache/prisma-python holds the npm install tree prisma-client-py + # drives; ~/.cache/prisma is where @prisma/engines stages its downloads. + path: | + ~/.cache/prisma-python + ~/.cache/prisma + key: ${{ runner.os }}-prisma-binaries-${{ steps.version.outputs.version }} diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index d34b0ee2e0f..e56f61988ef 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -83,7 +83,11 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac 🚄 Infrastructure ✅ Test -## Changes +## Caveats (if any) + + ## QA runbook diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index cee93bde7f2..58208988fca 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -18,10 +18,25 @@ on: type: number default: 2 timeout-minutes: - description: "Job timeout in minutes" + description: >- + Timeout for the test step alone. Setup (checkout, dependency install, + Prisma client generation) gets its own allowance on top, so a slow + runner or a cold binary download can never cancel passing tests. required: false type: number default: 20 + job-timeout-minutes: + description: >- + Backstop for the whole job. Keep it >= `timeout-minutes` plus 35: 30 for + the per-step ceilings on the setup steps below, and 5 for the runner + overhead the job clock charges but no step owns (job init, step + transitions, post-job cleanup). That headroom is what makes the test + budget a floor rather than a hope, since setup cannot overrun into it + without failing its own step first. GitHub expressions have no + arithmetic, so the sum is passed in rather than computed. + required: false + type: number + default: 55 max-failures: description: "Stop after this many failures" required: false @@ -44,30 +59,35 @@ jobs: run: name: Run tests runs-on: ubuntu-latest - timeout-minutes: ${{ inputs.timeout-minutes }} + timeout-minutes: ${{ inputs.job-timeout-minutes }} outputs: decision: ${{ steps.changes.outputs.decision }} steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + timeout-minutes: 3 with: persist-credentials: false - name: Detect backend-relevant changes id: changes + timeout-minutes: 2 uses: ./.github/actions/detect-backend-changes - name: Set up Python + timeout-minutes: 3 uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" - name: Set up uv + timeout-minutes: 3 uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" - name: Cache uv dependencies + timeout-minutes: 5 uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 with: path: | @@ -79,18 +99,24 @@ jobs: - name: Install dependencies if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 8 run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml + - name: Cache Prisma binaries + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 3 + uses: ./.github/actions/cache-prisma-binaries + - name: Generate Prisma client if: steps.changes.outputs.decision != 'skip' - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache + timeout-minutes: 3 run: | uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - name: Run tests if: steps.changes.outputs.decision != 'skip' + timeout-minutes: ${{ inputs.timeout-minutes }} env: TEST_PATH: ${{ inputs.test-path }} MAX_FAILURES: ${{ inputs.max-failures }} diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 02543d67a82..dbd663a2efa 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -71,10 +71,12 @@ jobs: if: steps.changes.outputs.relevant == 'true' run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + - name: Cache Prisma binaries + if: steps.changes.outputs.relevant == 'true' + uses: ./.github/actions/cache-prisma-binaries + - name: Generate Prisma client if: steps.changes.outputs.relevant == 'true' - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - name: Set up Node.js diff --git a/.github/workflows/mutation-test.yml b/.github/workflows/mutation-test.yml index da4fe073a6a..68317d5dd12 100644 --- a/.github/workflows/mutation-test.yml +++ b/.github/workflows/mutation-test.yml @@ -57,9 +57,10 @@ jobs: run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + - name: Generate Prisma client - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma diff --git a/.github/workflows/publish-basedpyright-base-counts.yml b/.github/workflows/publish-basedpyright-base-counts.yml index c85d30df0ce..71e196d8361 100644 --- a/.github/workflows/publish-basedpyright-base-counts.yml +++ b/.github/workflows/publish-basedpyright-base-counts.yml @@ -43,12 +43,13 @@ jobs: with: version: "0.10.9" + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + # The gate provisions its own measurement env (.venv-typecheck: a frozen # uv sync of its canonical dependency groups plus a generated Prisma # client), so no install step here can drift from what local runs measure. - name: Emit basedpyright counts for HEAD - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts" counts_file=$(ls "$RUNNER_TEMP"/basedpyright-counts/basedpyright-counts-*.json) diff --git a/.github/workflows/test-code-quality.yml b/.github/workflows/test-code-quality.yml index fab05fc2bbb..8f62837d29a 100644 --- a/.github/workflows/test-code-quality.yml +++ b/.github/workflows/test-code-quality.yml @@ -65,6 +65,12 @@ jobs: - name: check_provider_folders_documented run: uv run --no-sync python ./tests/code_coverage_tests/check_provider_folders_documented.py + - name: check_prisma_binary_cache + run: uv run --no-sync python ./tests/code_coverage_tests/check_prisma_binary_cache.py + + - name: check_workflow_startup_safety + run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py + - name: router_code_coverage run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 3db3fb07a94..280ec476cdf 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -71,12 +71,13 @@ jobs: run: | uv sync --frozen --group proxy-dev --group e2e-dev + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + # basedpyright resolves Prisma's generated client (litellm/proxy/schema.prisma) # only after `prisma generate` writes prisma/client.py et al. Without this the # DB wrappers typed against the generated client would degrade to Unknown. - name: Generate Prisma client - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma @@ -119,7 +120,6 @@ jobs: - name: Check basedpyright budget (delta vs base) env: GH_TOKEN: ${{ github.token }} - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA" diff --git a/.github/workflows/test-litellm-ui-unit.yml b/.github/workflows/test-litellm-ui-unit.yml index 8f2199017d9..69cbc082d98 100644 --- a/.github/workflows/test-litellm-ui-unit.yml +++ b/.github/workflows/test-litellm-ui-unit.yml @@ -42,6 +42,11 @@ jobs: - name: Install dependencies run: npm ci + - name: Run UI type tests (Vitest) + env: + CI: "true" + run: npm run test:types + - name: Run UI unit tests (Vitest) env: CI: "true" diff --git a/.github/workflows/test-terraform-provider.yml b/.github/workflows/test-terraform-provider.yml index 058a2538c15..7ea22825f4f 100644 --- a/.github/workflows/test-terraform-provider.yml +++ b/.github/workflows/test-terraform-provider.yml @@ -92,9 +92,10 @@ jobs: run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + - name: Generate Prisma client - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma diff --git a/.github/workflows/test-unit-documentation.yml b/.github/workflows/test-unit-documentation.yml index 50589cb5926..c93779c177f 100644 --- a/.github/workflows/test-unit-documentation.yml +++ b/.github/workflows/test-unit-documentation.yml @@ -65,10 +65,12 @@ jobs: run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + - name: Cache Prisma binaries + if: steps.changes.outputs.decision != 'skip' + uses: ./.github/actions/cache-prisma-binaries + - name: Generate Prisma client if: steps.changes.outputs.decision != 'skip' - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 60d2e471862..df212a85885 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -28,6 +28,10 @@ concurrency: # Most of a shard's time is pytest plugin load + xdist worker imports + # pytest-cov instrumentation, not the tests themselves. Keeping per-shard # work low and matching worker count to runner cores is what controls it. +# * `timeout` bounds the pytest step only. Checkout, dependency install, and +# Prisma client generation draw on a separate allowance in the base +# workflow, so slow setup shows up as a slow job rather than as a +# cancelled shard whose tests were passing. # * workers: 4 matches the 4-core ubuntu-latest runner. -n 8 on 4 cores # oversubscribes 2x and workers fight for CPU during their cold-start # imports (measured ~441% CPU for -n 8 locally, i.e. ~55% effective). diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index 2ea3c521e8b..64b92f7d847 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -76,4 +76,5 @@ jobs: workers: 4 reruns: 2 timeout-minutes: 60 + job-timeout-minutes: 95 artifact-name: proxy-server diff --git a/.github/workflows/test-unit-proxy-legacy.yml b/.github/workflows/test-unit-proxy-legacy.yml index 49aa5f9f51d..e8ca36fb30d 100644 --- a/.github/workflows/test-unit-proxy-legacy.yml +++ b/.github/workflows/test-unit-proxy-legacy.yml @@ -82,10 +82,12 @@ jobs: run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router + - name: Cache Prisma binaries + if: steps.changes.outputs.decision != 'skip' + uses: ./.github/actions/cache-prisma-binaries + - name: Generate Prisma client if: steps.changes.outputs.decision != 'skip' - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma diff --git a/.github/workflows/weekly_load_anomaly.yml b/.github/workflows/weekly_load_anomaly.yml index 4c2103f026d..2dffc889d0e 100644 --- a/.github/workflows/weekly_load_anomaly.yml +++ b/.github/workflows/weekly_load_anomaly.yml @@ -51,9 +51,10 @@ jobs: run: | .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy + - name: Cache Prisma binaries + uses: ./.github/actions/cache-prisma-binaries + - name: Generate Prisma client - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache run: | uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma diff --git a/CLAUDE.md b/CLAUDE.md index f1bb46c1fd3..436fa33fa41 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,4 +1,12 @@ -Do not write any comments (existing comments can stay) unless explicitly asked to in a user (not system) prompt +Do not write comments unless they are any of: +- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear) +- used as an input for tools to read and act on. For example: + - entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame + - a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # ` when introducing a truly unavoidable violation +- a TODO or FIXME + - Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work + +Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in: diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0385f7a96e7..96b689aed74 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,18 +1,18 @@ { "reportAny": { - "limit": 27731 + "limit": 26391 }, "reportArgumentType": { - "limit": 2626 + "limit": 2614 }, "reportAssignmentType": { - "limit": 329 + "limit": 327 }, "reportAttributeAccessIssue": { "limit": 514 }, "reportCallIssue": { - "limit": 116 + "limit": 114 }, "reportConstantRedefinition": { "limit": 40 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 8807 + "limit": 8319 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5835 + "limit": 5825 }, "reportMissingTypeArgument": { - "limit": 15790 + "limit": 15695 }, "reportMissingTypeStubs": { "limit": 40 @@ -90,7 +90,7 @@ "limit": 8 }, "reportReturnType": { - "limit": 217 + "limit": 213 }, "reportTypedDictNotRequiredAccess": { "limit": 26 @@ -99,22 +99,22 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45063 + "limit": 44996 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39773 + "limit": 39643 }, "reportUnknownParameterType": { - "limit": 20207 + "limit": 20132 }, "reportUnknownVariableType": { - "limit": 31281 + "limit": 31153 }, "reportUnnecessaryCast": { - "limit": 122 + "limit": 118 }, "reportUnnecessaryComparison": { "limit": 701 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 862 + "limit": 857 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 7acdd5dbdaf..dc8f17fb665 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,7 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t """ from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -43,11 +43,15 @@ class CheckBatchCost: # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True - async def _get_user_info(self, batch_id, user_id) -> dict: + async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]: """ Look up user email and key alias by user_id for enriching the S3 callback metadata. Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None). + Returns an empty dict when user_id is None: batches created by a team or service + account key carry no user id, and find_unique(where={"user_id": None}) raises. """ + if not user_id: + return {} try: user_row = await self.prisma_client.db.litellm_usertable.find_unique( where={"user_id": user_id} @@ -62,6 +66,66 @@ class CheckBatchCost: verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") return {} + async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None: + """Resolve the creating virtual key's alias from its hashed token.""" + if not api_key: + return None + try: + key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": api_key} + ) + return getattr(key_row, "key_alias", None) if key_row is not None else None + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}") + return None + + async def _get_team_alias(self, team_id: str | None) -> str | None: + """Resolve a team's alias from its id.""" + if not team_id: + return None + try: + team_row = await self.prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + return getattr(team_row, "team_alias", None) if team_row is not None else None + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}") + return None + + async def _build_creator_attribution_metadata( + self, job: "LiteLLM_ManagedObjectTable", batch_id: str + ) -> Dict[str, Any]: + """ + Rebuild the spend-tracking metadata for the key, team, and tags that created the + batch so the batch-cost spend log is attributed the same way a non-batch request + is. Rows created before api_key and request_tags were persisted carry only + created_by and team_id, and fall back to those. A named creating key owns + user_api_key_alias; when it has no alias, or the key has since been rotated or + deleted, the field keeps the creating user's alias that _get_user_info filled in, + because a resolvable name is more useful on the spend row than a null. + """ + api_key = getattr(job, "api_key", None) + team_id = getattr(job, "team_id", None) + request_tags = getattr(job, "request_tags", None) + + metadata: Dict[str, Any] = { + "user_api_key_user_id": job.created_by, + "user_api_key": api_key, + "user_api_key_team_id": team_id, + **(await self._get_user_info(batch_id, job.created_by)), + } + + key_alias = await self._get_key_alias(batch_id, api_key) + if key_alias is not None: + metadata["user_api_key_alias"] = key_alias + team_alias = await self._get_team_alias(team_id) + if team_alias is not None: + metadata["user_api_key_team_alias"] = team_alias + if isinstance(request_tags, list) and request_tags: + metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)] + + return metadata + async def _cleanup_stale_managed_objects(self) -> None: """ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days @@ -485,9 +549,6 @@ class CheckBatchCost: function_id=str(uuid.uuid4()), ) - creator_user_id = job.created_by - user_info = await self._get_user_info(batch_id, job.created_by) - logging_obj.update_environment_variables( litellm_params={ # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks @@ -496,11 +557,7 @@ class CheckBatchCost: "user-agent": CHECK_BATCH_COST_USER_AGENT, } }, - "metadata": { - "user_api_key_user_id": creator_user_id, - "user_api_key_team_id": getattr(job, "team_id", None), - **user_info, - }, + "metadata": await self._build_creator_attribution_metadata(job, batch_id), }, optional_params={}, ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index f0914240f79..37d267fcd6e 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -163,6 +163,8 @@ class _ManagedObjectTableActions(Protocol): self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] ) -> "PrismaManagedObjectRow": ... + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + class _CursorPageArgs(TypedDict, total=False): cursor: Mapping[str, str] @@ -263,7 +265,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_object_id: str, file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, + request_tags: Sequence[str] | None = None, + persist_attribution: bool = False, + create_if_missing: bool = True, ) -> None: + """Persist a managed object row, caching it and upserting it in the DB. + + persist_attribution is set only by the batch create, which is the one caller + that can speak for the creator; it gates the api_key and request_tags columns + that CheckBatchCost bills against, so a later poll or retrieve of the same + batch cannot record itself as the paying key. Like created_by and team_id, + both are written only in the upsert create branch, never on update. + + create_if_missing is cleared by callers that observe a batch they did not + create, such as a poll. They still refresh status and file_object, but a + row absent from the table is left absent rather than created with the + observer as its creator, because created_by and team_id are written from + whoever calls the create branch. + """ verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, @@ -277,6 +296,29 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span=litellm_parent_otel_span, ) + from prisma import Json + + api_key = user_api_key_dict.api_key or None + attribution_columns = ( + { + **({"api_key": api_key} if api_key is not None else {}), + **({"request_tags": Json(list(request_tags))} if request_tags else {}), + } + if persist_attribution + else {} + ) + # FIX: Update status and file_object on every operation to keep state in sync + update_columns: Final = { + "file_object": file_object.model_dump_json(), + "status": file_object.status, + "updated_by": user_api_key_dict.user_id, + } + if not create_if_missing: + await _managed_object_table(self.prisma_client).update_many( + where={"unified_object_id": unified_object_id}, + data=update_columns, + ) + return await _managed_object_table(self.prisma_client).upsert( where={"unified_object_id": unified_object_id}, data={ @@ -289,12 +331,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, "status": file_object.status, + **attribution_columns, }, - "update": { - "file_object": file_object.model_dump_json(), - "status": file_object.status, - "updated_by": user_api_key_dict.user_id, - }, # FIX: Update status and file_object on every operation to keep state in sync + "update": update_columns, }, ) @@ -1248,10 +1287,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): user_created_file_ids = await self.get_user_created_file_ids(user_api_key_dict, file_ids) ## Filter the response to only include the files created by the user response.data = user_created_file_ids # type: ignore + self._scope_list_page_cursors(response, user_created_file_ids) return response return response return response + @staticmethod + def _scope_list_page_cursors(response: AsyncCursorPage, data: List[OpenAIFileObject]) -> None: + """Rebuild ``first_id`` / ``last_id`` from the caller-scoped page. + + The upstream cursors point at rows that were just filtered out, so + leaving them in place discloses other callers' file ids. ``has_more`` + is always cleared because ``after`` is never forwarded upstream, so + no further page is reachable through the proxy. + """ + if hasattr(response, "first_id"): + response.first_id = data[0].id if data else None + if hasattr(response, "last_id"): + response.last_id = data[-1].id if data else None + if hasattr(response, "has_more"): + response.has_more = False + async def afile_retrieve( self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router: Optional[Router] = None ) -> OpenAIFileObject: diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713010000_add_ptu_columns_to_daily_team_spend/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713010000_add_ptu_columns_to_daily_team_spend/migration.sql new file mode 100644 index 00000000000..89a0494431b --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260713010000_add_ptu_columns_to_daily_team_spend/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "ptu_flat_cost" DOUBLE PRECISION NOT NULL DEFAULT 0.0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260730000000_add_api_key_and_request_tags_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260730000000_add_api_key_and_request_tags_to_managed_object_table/migration.sql new file mode 100644 index 00000000000..79bc6b24de8 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260730000000_add_api_key_and_request_tags_to_managed_object_table/migration.sql @@ -0,0 +1,5 @@ +-- Add api_key and request_tags columns to LiteLLM_ManagedObjectTable +-- Captured at batch-create time so CheckBatchCost can attribute batch-cost spend +-- back to the creating virtual key (and its tags) even when created_by is null. +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "api_key" TEXT; +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "request_tags" JSONB DEFAULT '[]'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9c871b65f40..33fd9389b63 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -30,7 +30,7 @@ model LiteLLM_BudgetTable { end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team - organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } // Models on proxy @@ -893,6 +893,7 @@ model LiteLLM_DailyTeamSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) created_at DateTime @default(now()) updated_at DateTime @updatedAt @@ -985,6 +986,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + api_key String? + request_tags Json? @default("[]") updated_at DateTime @updatedAt updated_by String? diff --git a/litellm/__init__.py b/litellm/__init__.py index e0c7d56361c..bc8a13ec2cd 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -197,6 +197,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = ( None # Fields to exclude from StandardLoggingPayload before callbacks receive it ) log_raw_request_response: bool = False +request_correlation_in_logs: bool = False redact_messages_in_exceptions: Optional[bool] = False redact_user_api_key_info: Optional[bool] = False # When True (default — preserves historical behavior), the Router appends diff --git a/litellm/_logging.py b/litellm/_logging.py index b9e102e2b3c..6add9d79a5b 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,4 +1,5 @@ import ast +import contextvars import logging import os import sys @@ -6,12 +7,44 @@ from datetime import datetime from logging import Formatter from typing import Any, Final +import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.secret_redaction import redact_string set_verbose = False +session_id_var: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("session_id", default="") +trace_id_var: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("trace_id", default="") + +_MAX_CORRELATION_ID_LENGTH: Final = 256 + + +def _sanitize_correlation_id(value: str) -> str: + """Strip control characters, bound length, and redact credential-shaped + content before a caller-controlled trace_id/session_id (e.g. + litellm_session_id, x-litellm-trace-id) is stamped into log lines. + + Without the first two, a caller could embed \\r/\\n or terminal escape + sequences to forge fake log entries, or submit an oversized value repeated + across every log line for the request. Without the redaction, a caller + could smuggle a real credential (e.g. an sk-... key) through this field: + CorrelationContextFilter stamps trace_id/session_id onto the record after + SecretRedactionFilter has already run, so those two fields never otherwise + pass through credential redaction. + """ + stripped: Final = "".join(ch for ch in value if ch.isprintable()) + return _redact_string(stripped[:_MAX_CORRELATION_ID_LENGTH]) + + +def set_session_id(session_id: str) -> "contextvars.Token[str]": + return session_id_var.set(_sanitize_correlation_id(session_id)) + + +def set_trace_id(trace_id: str) -> "contextvars.Token[str]": + return trace_id_var.set(_sanitize_correlation_id(trace_id)) + + if set_verbose is True: logging.warning( "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs." @@ -77,6 +110,28 @@ class SecretRedactionFilter(logging.Filter): _secret_filter: Final = SecretRedactionFilter() +class CorrelationContextFilter(logging.Filter): + """Stamps each log record with the current request's trace_id and session_id from contextvars. + + Works in tandem with JsonFormatter: the formatter's record.__dict__ loop picks up these + attributes as first-class JSON fields without any formatter-level code. + """ + + def filter(self, record: logging.LogRecord) -> bool: + if not litellm.request_correlation_in_logs: + return True + trace_id: Final = trace_id_var.get() + if trace_id: + record.trace_id = trace_id # rebind-ok: stamping the LogRecord is the Filter interface's contract + session_id: Final = session_id_var.get() + if session_id: + record.session_id = session_id # rebind-ok: stamping the LogRecord is the Filter interface's contract + return True + + +_correlation_filter: Final = CorrelationContextFilter() + + json_logs = bool(os.getenv("JSON_LOGS", False)) # Create a handler for the logger (you may need to adapt this based on your needs) log_level: Final = os.getenv("LITELLM_LOG", "DEBUG") @@ -84,6 +139,7 @@ numeric_level: Final[str] = getattr(logging, log_level.upper()) handler: Final = logging.StreamHandler() handler.setLevel(numeric_level) handler.addFilter(_secret_filter) +handler.addFilter(_correlation_filter) def _try_parse_json_message(message: str) -> dict[str, Any] | None: @@ -146,6 +202,11 @@ def _get_standard_record_attrs() -> frozenset: _STANDARD_RECORD_ATTRS: Final = _get_standard_record_attrs() +# CorrelationContextFilter is the only legitimate source for these two JSON fields; +# see JsonFormatter.format() for why they're excluded from the generic message-content +# and extra-attribute promotion paths. +_RESERVED_CORRELATION_FIELDS: Final = frozenset(("trace_id", "session_id")) + class JsonFormatter(Formatter): def __init__(self): @@ -164,13 +225,18 @@ class JsonFormatter(Formatter): "timestamp": self.formatTime(record), } - # Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties + # Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties. + # trace_id/session_id are excluded here unconditionally (not just "if not already + # set") - CorrelationContextFilter is the only legitimate source for these two + # fields, and a message that merely happens to parse as JSON/dict (e.g. a proxy + # log line dumping raw request headers) must never be able to claim them, even on + # a record the filter hasn't stamped yet (no correlation context active for it). parsed = _try_parse_json_message(message_str) if parsed is None: parsed = _try_parse_embedded_python_dict(message_str) if parsed is not None: for key, value in parsed.items(): - if key not in json_record: + if key not in json_record and key not in _RESERVED_CORRELATION_FIELDS: json_record[key] = value # Include extra attributes passed via logger.debug("msg", extra={...}) @@ -178,6 +244,18 @@ class JsonFormatter(Formatter): if key not in _STANDARD_RECORD_ATTRS and key not in json_record: json_record[key] = value + # trace_id/session_id are reserved: CorrelationContextFilter is the only + # legitimate source for these two fields. Without this, a message string + # that happens to parse as JSON/dict (e.g. a proxy log line dumping raw + # request headers) with a "trace_id"/"session_id" key would have already + # claimed the key at the parsed-message step above, and the extra-attributes + # loop's "key not in json_record" guard would then skip the real value - + # letting a caller-supplied header spoof another request's correlation ids. + for reserved_key in _RESERVED_CORRELATION_FIELDS: + value = getattr(record, reserved_key, None) + if value: + json_record[reserved_key] = value + # Set component/logger only if not already supplied via extra={...} if "component" not in json_record: json_record["component"] = record.name @@ -190,12 +268,34 @@ class JsonFormatter(Formatter): return safe_dumps(json_record) +class CorrelationPlainFormatter(logging.Formatter): + """Appends trace_id/session_id to plain-text log lines stamped by CorrelationContextFilter. + + Mirrors JsonFormatter's handling of these two fields so request_correlation_in_logs + behaves the same whether or not json_logs is enabled. + """ + + def format(self, record: logging.LogRecord) -> str: + formatted: Final = super().format(record) + trace_id: Final = getattr(record, "trace_id", None) + session_id: Final = getattr(record, "session_id", None) + if not trace_id and not session_id: + return formatted + parts: Final = tuple( + p + for p in (f"trace_id={trace_id}" if trace_id else None, f"session_id={session_id}" if session_id else None) + if p + ) + return f"{formatted} [{' '.join(parts)}]" + + # Function to set up exception handlers for JSON logging def _setup_json_exception_handlers(formatter): # Create a handler with JSON formatting for exceptions error_handler: Final = logging.StreamHandler() error_handler.setFormatter(formatter) error_handler.addFilter(_secret_filter) + error_handler.addFilter(_correlation_filter) # Setup excepthook for uncaught exceptions def json_excepthook(exc_type, exc_value, exc_traceback): @@ -243,7 +343,7 @@ if json_logs: handler.setFormatter(JsonFormatter()) _setup_json_exception_handlers(JsonFormatter()) else: - formatter: Final = logging.Formatter( + formatter: Final = CorrelationPlainFormatter( "\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s", datefmt="%H:%M:%S", ) @@ -346,6 +446,7 @@ def _initialize_loggers_with_handler(handler: logging.Handler): - Prevents bubbling to parent/root (critical to prevent duplicate JSON logs) """ handler.addFilter(_secret_filter) + handler.addFilter(_correlation_filter) for lg in _get_loggers_to_initialize(): lg.handlers.clear() # remove any existing handlers lg.addHandler(handler) # add JSON formatter handler diff --git a/litellm/constants.py b/litellm/constants.py index 3b91f23fe39..87d6fa1a744 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1493,6 +1493,8 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INT SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)) DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)) +RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500"))) +RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", "100"))) PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))) @@ -1719,3 +1721,18 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( ) UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS + +# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this +# sentinel api_key so PTU flat cost stays distinguishable from real per-request +# spend under the table's composite unique constraint. +PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__" +PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job" +PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900 +# Furthest back the catch-up pass looks for unpriced PTU days when a deployment +# declares no ptu_effective_from, bounding the scan for an open-ended window. +PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90 +# Slack allowed when deciding a sentinel row is stale. The row's updated_at and the +# run's cutoff are stamped by different hosts, so clock skew between them must not let +# one run delete a charge another just wrote. A stale row is hours old and a concurrent +# one is seconds old, so a few minutes separates them. +PTU_PRUNE_SKEW_GRACE_SECONDS: Final[int] = 300 diff --git a/litellm/evals/main.py b/litellm/evals/main.py index a25c7a96a8a..2f639d30ca0 100644 --- a/litellm/evals/main.py +++ b/litellm/evals/main.py @@ -7,7 +7,7 @@ import asyncio import contextvars from collections.abc import Coroutine from functools import partial -from typing import Any, Final +from typing import Final import httpx @@ -21,8 +21,10 @@ from litellm.types.llms.openai_evals import ( CancelRunResponse, CreateEvalRequest, CreateRunRequest, + DataSourceConfig, DeleteEvalResponse, Eval, + GraderConfig, ListEvalsParams, ListEvalsResponse, ListRunsParams, @@ -41,13 +43,13 @@ DEFAULT_OPENAI_API_BASE: Final = "https://api.openai.com" @client async def acreate_eval( - data_source_config: dict[str, Any], - testing_criteria: list[dict[str, Any]], + data_source_config: DataSourceConfig, + testing_criteria: list[GraderConfig], name: str | None = None, - metadata: dict[str, Any] | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -110,17 +112,17 @@ async def acreate_eval( @client def create_eval( - data_source_config: dict[str, Any], - testing_criteria: list[dict[str, Any]], + data_source_config: DataSourceConfig, + testing_criteria: list[GraderConfig], name: str | None = None, - metadata: dict[str, Any] | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> Eval | Coroutine[Any, Any, Eval]: +) -> Eval | Coroutine[object, object, Eval]: """ Create a new evaluation @@ -231,8 +233,8 @@ async def alist_evals( before: str | None = None, order: str | None = None, order_by: str | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -300,12 +302,12 @@ def list_evals( before: str | None = None, order: str | None = None, order_by: str | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> ListEvalsResponse | Coroutine[Any, Any, ListEvalsResponse]: +) -> ListEvalsResponse | Coroutine[object, object, ListEvalsResponse]: """ List all evaluations @@ -413,8 +415,8 @@ def list_evals( @client async def aget_eval( eval_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -470,12 +472,12 @@ async def aget_eval( @client def get_eval( eval_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> Eval | Coroutine[Any, Any, Eval]: +) -> Eval | Coroutine[object, object, Eval]: """ Get an evaluation by ID @@ -564,10 +566,10 @@ def get_eval( async def aupdate_eval( eval_id: str, name: str | None = None, - metadata: dict[str, Any] | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -630,14 +632,14 @@ async def aupdate_eval( def update_eval( eval_id: str, name: str | None = None, - metadata: dict[str, Any] | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> Eval | Coroutine[Any, Any, Eval]: +) -> Eval | Coroutine[object, object, Eval]: """ Update an evaluation @@ -783,8 +785,8 @@ def update_eval( @client async def adelete_eval( eval_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -840,12 +842,12 @@ async def adelete_eval( @client def delete_eval( eval_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> DeleteEvalResponse | Coroutine[Any, Any, DeleteEvalResponse]: +) -> DeleteEvalResponse | Coroutine[object, object, DeleteEvalResponse]: """ Delete an evaluation @@ -933,8 +935,8 @@ def delete_eval( @client async def acancel_eval( eval_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -990,12 +992,12 @@ async def acancel_eval( @client def cancel_eval( eval_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> CancelEvalResponse | Coroutine[Any, Any, CancelEvalResponse]: +) -> CancelEvalResponse | Coroutine[object, object, CancelEvalResponse]: """ Cancel a running evaluation @@ -1092,12 +1094,12 @@ def cancel_eval( @client async def acreate_run( eval_id: str, - data_source: dict[str, Any], + data_source: dict[str, object], name: str | None = None, - metadata: dict[str, Any] | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -1161,16 +1163,16 @@ async def acreate_run( @client def create_run( eval_id: str, - data_source: dict[str, Any], + data_source: dict[str, object], name: str | None = None, - metadata: dict[str, Any] | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> Run | Coroutine[Any, Any, Run]: +) -> Run | Coroutine[object, object, Run]: """ Create a new run for an evaluation @@ -1280,8 +1282,8 @@ async def alist_runs( after: str | None = None, before: str | None = None, order: str | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -1349,12 +1351,12 @@ def list_runs( after: str | None = None, before: str | None = None, order: str | None = None, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> ListRunsResponse | Coroutine[Any, Any, ListRunsResponse]: +) -> ListRunsResponse | Coroutine[object, object, ListRunsResponse]: """ List all runs for an evaluation @@ -1462,8 +1464,8 @@ def list_runs( async def aget_run( eval_id: str, run_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -1522,12 +1524,12 @@ async def aget_run( def get_run( eval_id: str, run_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> Run | Coroutine[Any, Any, Run]: +) -> Run | Coroutine[object, object, Run]: """ Get a specific run @@ -1618,8 +1620,8 @@ def get_run( async def acancel_run( eval_id: str, run_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -1678,12 +1680,12 @@ async def acancel_run( def cancel_run( eval_id: str, run_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> CancelRunResponse | Coroutine[Any, Any, CancelRunResponse]: +) -> CancelRunResponse | Coroutine[object, object, CancelRunResponse]: """ Cancel a running run @@ -1783,8 +1785,8 @@ def cancel_run( async def adelete_run( eval_id: str, run_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -1843,12 +1845,12 @@ async def adelete_run( def delete_run( eval_id: str, run_id: str, - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> RunDeleteResponse | Coroutine[Any, Any, RunDeleteResponse]: +) -> RunDeleteResponse | Coroutine[object, object, RunDeleteResponse]: """ Delete a run diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 8c494794858..e7e1ab538d5 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -1,4 +1,5 @@ import json +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from typing_extensions import override @@ -12,7 +13,7 @@ from litellm.litellm_core_utils.redact_messages import ( should_redact_message_logging, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall, StandardLoggingPayload if TYPE_CHECKING: from opentelemetry.trace import Span @@ -22,6 +23,7 @@ from litellm.integrations._types.open_inference import ( ImageAttributes, MessageAttributes, MessageContentAttributes, + OpenInferenceMimeTypeValues, OpenInferenceSpanKindValues, SpanAttributes, ToolCallAttributes, @@ -480,6 +482,7 @@ def set_attributes(span: "Span", kwargs, response_obj, attributes: type[BaseLLMO response_obj_for_attrs, slp, ) + _safe_emit("mcp tool attrs", _maybe_set_mcp_tool_attrs, span, kwargs, slp, response_obj_for_attrs) def _sanitize_optional_params(optional_params: dict | None) -> dict: @@ -538,9 +541,12 @@ def _set_request_attributes( if optional_params.get("user"): safe_set_attribute(span, "llm.user", optional_params.get("user")) - if response_obj and response_obj.get("id"): + if not hasattr(response_obj, "get"): + return + + if response_obj.get("id"): safe_set_attribute(span, "llm.response.id", response_obj.get("id")) - if response_obj and response_obj.get("model"): + if response_obj.get("model"): safe_set_attribute(span, "llm.response.model", response_obj.get("model")) @@ -588,6 +594,8 @@ def _coerce_response_obj_for_attrs(response_obj): - dicts and Pydantic models that already expose `.get` are returned unchanged (preserves all current behavior, including the Responses API flow which relies on Pydantic attribute access). + - Pydantic models without `.get` (e.g. the MCP SDK's `CallToolResult`, + logged for `call_mcp_tool` spans) are dumped to a dict. - `httpx.Response` and other text-only responses (passthrough routes) are JSON-decoded so the standard extraction paths can read fields like `id`, `model`, and `usage`. On failure the original object is returned @@ -595,6 +603,9 @@ def _coerce_response_obj_for_attrs(response_obj): """ if response_obj is None or hasattr(response_obj, "get"): return response_obj + dumped: Final = _to_plain_dict(response_obj) + if isinstance(dumped, dict): + return dumped text: Final = getattr(response_obj, "text", None) if isinstance(text, str) and text: try: @@ -1058,3 +1069,65 @@ def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs): except Exception: return None return None + + +def _maybe_set_mcp_tool_attrs( + span: "Span", + kwargs: Mapping[str, object], + standard_logging_payload: StandardLoggingPayload | None, + coerced_response_obj: object, +) -> None: + """Render `call_mcp_tool` spans as OpenInference TOOL spans. + + MCP tool calls carry neither `messages` nor `choices`, so the generic + extraction paths leave Input/Output blank. The tool name and arguments live + in `metadata.mcp_tool_call_metadata`; the result is an MCP `CallToolResult` + whose `content` is a list of typed parts. + """ + if standard_logging_payload is None: + return + if standard_logging_payload.get("call_type") != CallTypes.call_mcp_tool.value: + return + + metadata: Final = standard_logging_payload.get("metadata") + mcp_meta: Final[StandardLoggingMCPToolCall | None] = metadata.get("mcp_tool_call_metadata") if metadata else None + if mcp_meta is None: + return + + tool_name: Final = mcp_meta.get("name") or mcp_meta.get("namespaced_tool_name") + if tool_name: + safe_set_attribute(span, SpanAttributes.TOOL_NAME, tool_name) + + if should_redact_message_logging(kwargs): # pyright: ignore[reportArgumentType] # reads, never mutates + return + + arguments: Final[object] = mcp_meta.get("arguments") + if arguments is not None: + safe_set_attribute(span, SpanAttributes.INPUT_VALUE, safe_dumps(arguments)) + safe_set_attribute(span, SpanAttributes.INPUT_MIME_TYPE, OpenInferenceMimeTypeValues.JSON.value) + + _set_mcp_tool_output(span, coerced_response_obj) + + +def _has_only_text_parts(content: object) -> bool: + return not isinstance(content, list) or all(_coerce_text([part]) is not None for part in content) + + +def _set_mcp_tool_output(span: "Span", coerced_response_obj: object) -> None: + if not isinstance(coerced_response_obj, Mapping): + return + + content: Final[object] = coerced_response_obj.get("content") + text: Final[str | None] = _coerce_text(content) + if text and _has_only_text_parts(content): + safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, text) + safe_set_attribute(span, SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.TEXT.value) + return + + structured: Final[object] = coerced_response_obj.get("structuredContent") + payload: Final[object] = content if content else structured if structured is not None else content + if payload is None: + return + + safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, safe_dumps(payload)) + safe_set_attribute(span, SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.JSON.value) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 2c5f7484dac..f7f27459768 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,7 +10,7 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -73,6 +73,18 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" +class _PlanMetadataView(TypedDict): + websearch_native_blocks: Sequence[Mapping[str, object]] | None + + +class _AgenticLoopParamsView(TypedDict): + agentic_loop_params: AgenticLoopParams + + +class _WebSearchSettingsView(TypedDict): + websearch_interception_params: WebSearchInterceptionConfig + + class WebSearchInterceptionLogger(CustomLogger): """ CustomLogger that intercepts WebSearch tool calls for models that don't @@ -394,7 +406,7 @@ class WebSearchInterceptionLogger(CustomLogger): return tool.get("name") @classmethod - def _sync_forced_tool_choice(cls, tool_choice: Any, converted_tools: list[dict[str, object]]) -> object: + def _sync_forced_tool_choice(cls, tool_choice: object, converted_tools: Sequence[Mapping[str, object]]) -> object: """Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it names a web-search tool that was just converted away. @@ -462,7 +474,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs[WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY] = True # Convert native web search tools to LiteLLM standard - converted_tools: Final = [] + converted_tools: Final[list[dict[str, object]]] = [] for tool in tools: if is_web_search_tool(tool): standard_tool = get_litellm_web_search_tool() @@ -833,7 +845,10 @@ class WebSearchInterceptionLogger(CustomLogger): Anthropic-native clients (Claude Desktop, the Anthropic SDK) can render citations / sources alongside the model's textual reply. """ - native_blocks: Final = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) + metadata_view: Final[_PlanMetadataView] = { + "websearch_native_blocks": plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) + } + native_blocks: Final = metadata_view["websearch_native_blocks"] if not native_blocks: return response return self._inject_native_blocks(response, native_blocks) @@ -1278,8 +1293,10 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: - agentic_params: Final[AgenticLoopParams] = logging_obj.model_call_details.get("agentic_loop_params", {}) - full_model_name = agentic_params.get("model", model) + agentic_view: Final[_AgenticLoopParamsView] = { + "agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {}) + } + full_model_name = agentic_view["agentic_loop_params"].get("model", model) verbose_logger.debug( "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]", _call_id, @@ -1676,7 +1693,7 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( litellm_settings: dict[str, Any], - callback_specific_params: dict[str, Any], + callback_specific_params: Mapping[str, object], ) -> "WebSearchInterceptionLogger": """ Static method to initialize WebSearchInterceptionLogger from proxy config. @@ -1700,7 +1717,10 @@ class WebSearchInterceptionLogger(CustomLogger): # Get websearch_interception_params from litellm_settings or callback_specific_params websearch_params: WebSearchInterceptionConfig = {} if "websearch_interception_params" in litellm_settings: - websearch_params = litellm_settings["websearch_interception_params"] + settings_view: Final[_WebSearchSettingsView] = { + "websearch_interception_params": litellm_settings["websearch_interception_params"] + } + websearch_params = settings_view["websearch_interception_params"] elif "websearch_interception" in callback_specific_params and isinstance( callback_specific_params["websearch_interception"], dict ): diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index fe986dd5fce..a44ce431f4e 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -8,6 +8,7 @@ This module has no dependencies on proxy code and can be safely imported at the import json import os import time +from collections.abc import Mapping from pathlib import Path from typing import Final @@ -71,7 +72,7 @@ def get_litellm_gateway_api_key( return token_data["key"] -def is_cli_token_fresh(token_data: dict, buffer_hours: float = 0.1) -> bool: +def is_cli_token_fresh(token_data: Mapping[str, object], buffer_hours: float = 0.1) -> bool: """Check whether a cached CLI token (as stored in token.json) is still within its expiration window. Used by `lite auth print-token` to fail fast, without a network round trip, once the cached token is past diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0e115946bc7..1da03c1db4e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -10,7 +10,7 @@ import subprocess import sys import time import traceback -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime as dt_object from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast @@ -25,7 +25,15 @@ from litellm import ( log_raw_request_response, turn_off_message_logging, ) -from litellm._logging import _is_debugging_on, _redact_string, verbose_logger +from litellm._logging import ( + _is_debugging_on, + _redact_string, + session_id_var, + set_session_id, + set_trace_id, + trace_id_var, + verbose_logger, +) from litellm._uuid import uuid from litellm.batches.batch_utils import _handle_completed_batch from litellm.caching.caching import DualCache, InMemoryCache @@ -313,6 +321,7 @@ class Logging(LiteLLMLoggingBaseClass): applied_guardrails: list[str] | None = None, kwargs: dict | None = None, log_raw_request_response: bool = False, + supports_correlation_logging: bool = True, ): _input: Final[str | None] = messages # save original value of messages if messages is not None: @@ -338,6 +347,36 @@ class Logging(LiteLLMLoggingBaseClass): self.call_type = call_type self.litellm_call_id = litellm_call_id self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4()) + + # Capture the pre-call *value* (not a contextvars.Token) so restoration works + # even if this attempt's own logging ends up dispatched onto a different + # asyncio Task/context (e.g. via asyncio.create_task or the logging worker) - + # a Token can only be reset in the exact Context where it was created. + self._pre_call_trace_id: str = trace_id_var.get() + self._pre_call_session_id: str = session_id_var.get() + _sid: Final = kwargs.get("litellm_session_id") if kwargs else None + self.litellm_session_id: str = str(_sid) if _sid else "" + # supports_correlation_logging is False for calls originating from the + # sync client entry point (wrapper() in utils.py): a plain OS thread + # has no per-call context isolation the way an asyncio Task does, and + # a thread pool's worker threads are recycled across unrelated + # requests, so stamping trace_id/session_id there risks one request's + # ids leaking into a different, later request on the same thread. Sync + # support is deferred to a follow-up PR with its own safe-restore + # mechanism; async calls (the proxy's only call path) are unaffected. + if supports_correlation_logging: + set_trace_id(self.litellm_trace_id) + set_session_id(self.litellm_session_id) + # set_trace_id()/set_session_id() sanitize (strip control chars, bound + # length) before storing, so the contextvar's actual value can differ + # from self.litellm_trace_id/litellm_session_id. Capture what was + # really stored - _restore_correlation_context_if_unclaimed() must + # compare against this, not the raw ids, or a caller-supplied id + # containing control characters/oversized input would never match + # and cleanup would be skipped forever. + self._own_trace_id: str = trace_id_var.get() + self._own_session_id: str = session_id_var.get() + self.function_id = function_id self.streaming_chunks: list[Any] = [] # for generating complete stream response self.sync_streaming_chunks: list[Any] = [] # for generating complete stream response @@ -1992,7 +2031,67 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: await self.async_success_handler(result=complete_streaming_response) - def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): + def _restore_correlation_context(self) -> None: + """Restore trace_id/session_id contextvars to their pre-call value. + + Without this, a nested LiteLLM call sharing the same asyncio Task as an + outer request (e.g. a guardrail's own LLM-as-judge call, an MCP sampling + call) would leave the outer request's subsequent log lines stamped with + the nested call's trace_id/session_id instead of its own. + + Uses a plain set() of the captured pre-call value rather than + contextvars.Token-based reset(), since this can end up called from a + different asyncio Task/context than __init__ ran in (e.g. the request + task's own wrapper() finally block, plus async_success_handler + dispatched separately via asyncio.create_task/the logging worker) - + reset() only works in the exact Context a Token was created in and + raises otherwise. Deliberately NOT idempotent/guarded: each distinct + Task that calls this needs its own restore to actually take effect in + that Task's view of the contextvars, so calling it multiple times + (once per Task involved in this attempt) is required, not just safe. + """ + set_trace_id(self._pre_call_trace_id) + set_session_id(self._pre_call_session_id) + + def _restore_correlation_context_if_unclaimed(self) -> None: + """Guarded variant for __del__-triggered cleanup only. + + __del__ can fire arbitrarily late (delayed by cyclic GC, possibly + after the consuming Task/thread has already moved on to a different, + still-active call). Unconditionally restoring in that case would + stomp the active call's trace_id/session_id with this abandoned + stream's stale pre-call snapshot. Only restore if the contextvars + still hold the ids *this* call set - i.e. nothing has claimed them + since - so an unrelated active call is never overwritten. + """ + if trace_id_var.get() == self._own_trace_id and session_id_var.get() == self._own_session_id: + self._restore_correlation_context() + + def success_handler( + self, + result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: Any, # kwargs-ok: forwarded to _success_handler_body + ) -> None: + """Restores trace_id/session_id contextvars once this attempt's own success + logging (including any nested calls its callbacks trigger) is fully done.""" + try: + return self._success_handler_body( + result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs + ) + finally: + self._restore_correlation_context() + + def _success_handler_body( + self, + result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: Any, # kwargs-ok: forwarded from success_handler + ) -> None: verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit) if not self.should_run_logging(event_type="sync_success"): # prevent double logging return @@ -2399,7 +2498,31 @@ class Logging(LiteLLMLoggingBaseClass): e, ) - async def async_success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs): + async def async_success_handler( + self, + result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: Any, # kwargs-ok: forwarded to _async_success_handler_body + ) -> None: + """Restores trace_id/session_id contextvars once this attempt's own success + logging (including any nested calls its callbacks trigger) is fully done.""" + try: + return await self._async_success_handler_body( + result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs + ) + finally: + self._restore_correlation_context() + + async def _async_success_handler_body( + self, + result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: Any, # kwargs-ok: forwarded from async_success_handler + ) -> None: """ Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. """ @@ -2791,7 +2914,32 @@ class Logging(LiteLLMLoggingBaseClass): kwargs=self.model_call_details, ) - def failure_handler(self, exception, traceback_exception, start_time=None, end_time=None): + def failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: + """Restores trace_id/session_id contextvars once this attempt's own failure + logging (including any nested calls its callbacks trigger) is fully done.""" + try: + return self._failure_handler_body( + exception=exception, + traceback_exception=traceback_exception, + start_time=start_time, + end_time=end_time, + ) + finally: + self._restore_correlation_context() + + def _failure_handler_body( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback) if not self.should_run_logging(event_type="sync_failure"): # prevent double logging return @@ -2960,7 +3108,32 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging %s", e ) - async def async_failure_handler(self, exception, traceback_exception, start_time=None, end_time=None): + async def async_failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: + """Restores trace_id/session_id contextvars once this attempt's own failure + logging (including any nested calls its callbacks trigger) is fully done.""" + try: + return await self._async_failure_handler_body( + exception=exception, + traceback_exception=traceback_exception, + start_time=start_time, + end_time=end_time, + ) + finally: + self._restore_correlation_context() + + async def _async_failure_handler_body( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: """ Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions. """ @@ -5074,33 +5247,61 @@ class StandardLoggingPayloadSetup: return end_time_float - start_time_float @staticmethod - def _get_standard_logging_payload_trace_id( + def get_standard_logging_payload_trace_id( logging_obj: Logging, - litellm_params: dict, + litellm_params: Mapping[str, Any], ) -> str: """ Returns the `litellm_trace_id` for this request This helps link sessions when multiple requests are made in a single session + + Gated behind `litellm.request_correlation_in_logs`: + - Off (default): legacy behavior, preserved for backward compatibility - + `litellm_session_id` takes priority over `litellm_trace_id` since historically + this field doubled as the session-grouping field. + - On: `litellm_trace_id` takes priority - trace_id and session_id are independent, + see `get_standard_logging_payload_session_id` for session tracking. """ dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id") dynamic_litellm_trace_id: Final = litellm_params.get("litellm_trace_id") + metadata: Final = litellm_params.get("metadata") + metadata_session_id: Final = metadata.get("session_id") if metadata else None + metadata_trace_id: Final = metadata.get("trace_id") if metadata else None - # Note: we recommend using `litellm_session_id` for session tracking - # `litellm_trace_id` is an internal litellm param + ordered_candidates: Final[tuple[Any, Any, Any, Any]] = ( + (dynamic_litellm_trace_id, dynamic_litellm_session_id, metadata_trace_id, metadata_session_id) + if litellm.request_correlation_in_logs + else (dynamic_litellm_session_id, dynamic_litellm_trace_id, metadata_session_id, metadata_trace_id) + ) + for candidate in ordered_candidates: + if candidate: + return str(candidate) + return logging_obj.litellm_trace_id + + @staticmethod + def get_standard_logging_payload_session_id( + logging_obj: Logging, + litellm_params: Mapping[str, Any], + ) -> str: + """ + Returns the end-user/conversation `litellm_session_id` for this request, independent of trace_id. + + Only populated when `litellm.request_correlation_in_logs` is enabled - off by default + to avoid changing existing StandardLoggingPayload shape for callers who haven't opted in. + Unlike `get_standard_logging_payload_trace_id`, this never falls back to a generated + per-call trace id: it's empty when the caller never supplied a session id. + """ + if not litellm.request_correlation_in_logs: + return "" + dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id") if dynamic_litellm_session_id: return str(dynamic_litellm_session_id) - elif dynamic_litellm_trace_id: - return str(dynamic_litellm_trace_id) - # Fallback: use metadata.session_id or metadata.trace_id for call chaining - metadata: Final = litellm_params.get("metadata") or {} - metadata_session_id: Final = metadata.get("session_id") - metadata_trace_id: Final = metadata.get("trace_id") + metadata: Final = litellm_params.get("metadata") + metadata_session_id: Final = metadata.get("session_id") if metadata else None if metadata_session_id: return str(metadata_session_id) - if metadata_trace_id: - return str(metadata_trace_id) - return logging_obj.litellm_trace_id + return logging_obj.litellm_session_id @staticmethod def _get_user_agent_tags(proxy_server_request: dict) -> list[str] | None: @@ -5405,7 +5606,11 @@ def get_standard_logging_object_payload( payload: Final[StandardLoggingPayload] = StandardLoggingPayload( id=str(id), litellm_call_id=kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), - trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + trace_id=StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=logging_obj, + litellm_params=litellm_params, + ), + session_id=StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( logging_obj=logging_obj, litellm_params=litellm_params, ), diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 6574b261774..b94851794f0 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -49,6 +49,12 @@ _SERVICE_TIER_TO_COST_KEY_SUFFIX: Final[Mapping[str, str]] = MappingProxyType( } ) +_INCLUSIVE_THRESHOLD_PROVIDERS: Final = frozenset({"xai"}) + + +def _uses_inclusive_token_thresholds(custom_llm_provider: str | None) -> bool: + return custom_llm_provider in _INCLUSIVE_THRESHOLD_PROVIDERS + def _get_token_detail_value(details: object, key: str) -> int | None: if isinstance(details, dict): @@ -202,7 +208,11 @@ def _parse_above_token_threshold(key: str) -> float: def _get_token_base_cost( - model_info: ModelInfo, usage: Usage, service_tier: str | None = None + model_info: ModelInfo, + usage: Usage, + service_tier: str | None = None, + *, + threshold_is_inclusive: bool = False, ) -> tuple[float, float, float, float, float]: """ Return prompt cost, completion cost, and cache costs for a given model and usage. @@ -210,6 +220,9 @@ def _get_token_base_cost( If input_tokens > threshold and `input_cost_per_token_above_[x]k_tokens` or `input_cost_per_token_above_[x]_tokens` is set, then we use the corresponding threshold cost for all token types. + `threshold_is_inclusive` switches that comparison to >=, for providers such as xAI + that bill the higher tier once the prompt reaches the threshold. + Returns: Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost) """ @@ -262,7 +275,7 @@ def _get_token_base_cost( # Handle both formats: _above_128k_tokens and _above_128_tokens threshold_str = key.split("_above_")[1].split("_tokens")[0] threshold = _parse_above_token_threshold(key) - if usage.prompt_tokens > threshold: + if usage.prompt_tokens > threshold or (threshold_is_inclusive and usage.prompt_tokens == threshold): # Prefer a service_tier-specific above-threshold key when available, # e.g. input_cost_per_token_priority_above_200k_tokens for Gemini # ON_DEMAND_PRIORITY. Falls back to the standard key automatically @@ -777,7 +790,12 @@ def generic_cost_per_token( 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) + ) = _get_token_base_cost( + model_info=model_info, + usage=usage, + service_tier=service_tier, + threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), + ) prompt_cost = _calculate_input_cost( prompt_tokens_details=prompt_tokens_details, @@ -909,7 +927,12 @@ def get_token_type_cost_breakdown( cache_creation_cost_rate, cache_creation_cost_above_1hr_rate, cache_read_cost_rate, - ) = _get_token_base_cost(model_info=model_info, usage=usage, service_tier=service_tier) + ) = _get_token_base_cost( + model_info=model_info, + usage=usage, + service_tier=service_tier, + threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), + ) reasoning_tokens = ( _parse_completion_tokens_details(usage)["reasoning_tokens"] @@ -996,9 +1019,13 @@ def calculate_image_response_cost_from_usage( input_tokens_details: Final = getattr(usage, "input_tokens_details", None) prompt_tokens_details: PromptTokensDetailsWrapper | None = None if input_tokens_details is not None: + # input_tokens_details may be a dict (e.g. OpenAI image edit responses) + # or an object; read it tolerantly like the output side below, so image + # input tokens are priced at input_cost_per_image_token instead of + # silently falling back to the text rate. prompt_tokens_details = PromptTokensDetailsWrapper( - text_tokens=getattr(input_tokens_details, "text_tokens", None), - image_tokens=getattr(input_tokens_details, "image_tokens", None), + text_tokens=_get_token_detail_value(input_tokens_details, "text_tokens"), + image_tokens=_get_token_detail_value(input_tokens_details, "image_tokens"), cached_tokens=0, ) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 68465d06b15..2dc71abee3e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -213,7 +213,75 @@ class CustomStreamWrapper: def __aiter__(self) -> AsyncIterator["ModelResponseStream"]: return self + def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None: + """Restore trace_id/session_id in the *consuming* thread/task/context. + + wrapper_async() deliberately skips restoring correlation context when + it returns a stream, so log lines emitted while the caller iterates it + still carry this call's ids (see request_correlation_in_logs). + wrapper() (the sync path) never stamps anything in the first place - + see Logging.__init__'s supports_correlation_logging - so this method + is an inert no-op for sync-created streams, harmless to call anyway + since the class is shared between __next__ and __anext__. + But the terminal success/failure handlers this stream dispatches to + finish the job run on a *different* Task/thread (asyncio.create_task, + threading.Thread, or the shared executor) - restoring there fixes up + that detached context, not the one actually running the caller's + `for`/`async for` loop. Call this at every point control genuinely + returns to that consuming context: natural exhaustion (StopIteration/ + StopAsyncIteration), a raised failure, or explicit aclose(). Never let + this raise - it must not break the caller's actual stream handling. + + guarded=True (only __del__ uses this) skips the restore unless the + contextvars still hold the ids this stream's own call set, so a + delayed finalizer never overwrites a different, still-active call + that has since taken over the same Task/thread's context. + """ + try: + logging_obj: Final = getattr(self, "logging_obj", None) + if logging_obj is None: + return + method_name: Final = ( + "_restore_correlation_context_if_unclaimed" if guarded else "_restore_correlation_context" + ) + restore: Final = getattr(logging_obj, method_name, None) + if restore is not None: + restore() + except Exception as restore_error: # noqa: BLE001 # best-effort cleanup; must not raise into the caller + verbose_logger.debug("could not restore correlation context: %s", restore_error) + + def __del__(self) -> None: + """Best-effort correlation-context cleanup for an abandoned async stream. + + Only meaningfully applies to streams created by wrapper_async(): it + leaves contextvars "open" across the caller's iteration, so if the + caller never fully consumes the stream - stops early, drops the + reference, cancels it - none of the exit points + _restore_consumer_correlation_context() is called from ever run. For a + sync stream (wrapper()), this is a no-op in practice: wrapper() never + stamps trace_id/session_id for sync calls in the first place (see + Logging.__init__'s supports_correlation_logging), so there is nothing + for this to clean up. + + This is a best-effort fallback, not a guarantee: __del__ timing is + unpredictable (delayed by cyclic GC, not guaranteed at interpreter + shutdown, and may run on a different thread), so this can only reduce + how long the leak persists, not eliminate it. That's an acceptable + trade specifically because its blast radius is bounded to the one + asyncio Task this stream's own call ran in - each async call has its + own copy of the contextvars, and Tasks (unlike a thread pool's worker + threads) are never recycled across requests, so a delayed or missed + cleanup here can never misattribute a *different* request's logs. + guarded=True additionally ensures it never clobbers a different, + still-active call's context within that same Task if this fires late. + """ + self._restore_consumer_correlation_context(guarded=True) + async def aclose(self): + # Restore the consumer's outer context only after the underlying + # provider stream's own close (and its diagnostic logging below, if + # closing fails) completes - not before - so those log lines still + # carry this closing stream's own trace_id/session_id. if self.completion_stream is not None: stream_to_close: Final = self.completion_stream self.completion_stream = None @@ -233,6 +301,7 @@ class CustomStreamWrapper: "CustomStreamWrapper.aclose: error closing completion_stream: %s", e, ) + self._restore_consumer_correlation_context() def check_send_stream_usage(self, stream_options: dict | None): return stream_options is not None and stream_options.get("include_usage", False) is True @@ -1839,6 +1908,7 @@ class CustomStreamWrapper: if self.sent_stream_usage is False and self.send_stream_usage is True: self.sent_stream_usage = True return response + self._restore_consumer_correlation_context() raise # Re-raise StopIteration else: self.sent_last_chunk = True @@ -1852,6 +1922,19 @@ class CustomStreamWrapper: processed_chunk, cache_hit, ) # log response + # Deliberately do NOT restore context here even though + # completion_stream is already exhausted: this chunk is still + # real data belonging to this call, and the caller's own + # (application-level) log statements processing it run + # immediately after this return, in this same synchronous + # frame - restoring first would make those lines carry the + # wrong ids, which is exactly what leaving context open during + # iteration is meant to prevent (see + # _restore_consumer_correlation_context's docstring). A caller + # that keeps iterating gets cleaned up on its next __next__() + # call (immediate StopIteration, handled above); one that + # stops right here relies on aclose() or the best-effort + # __del__ guard instead. return processed_chunk except Exception as e: traceback_exception: Final = traceback.format_exc() @@ -1879,8 +1962,12 @@ class CustomStreamWrapper: cache_hit = False if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": cache_hit = True - self._check_max_streaming_duration() try: + # Inside the try (not before it) so a raised litellm.Timeout flows + # through the same except Exception -> _handle_stream_fallback_error + # path as every other failure, restoring the consumer's correlation + # context - a check before the try would bypass that entirely. + self._check_max_streaming_duration() if self.completion_stream is None: await self.fetch_stream() @@ -2083,10 +2170,17 @@ class CustomStreamWrapper: ) ) + self._restore_consumer_correlation_context() raise StopAsyncIteration # Re-raise StopIteration else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() + # see sync __next__'s sibling branch: deliberately do NOT restore + # here - this chunk is still this call's own data, and restoring + # before returning it would corrupt the caller's own log + # statements processing it. A caller that keeps iterating gets + # cleaned up on the next __anext__() call; one that stops here + # relies on aclose() or the best-effort __del__ guard. return processed_chunk def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn: @@ -2138,7 +2232,12 @@ class CustomStreamWrapper: """ from litellm.exceptions import MidStreamFallbackError - # Map to OpenAI exception format + # Map to OpenAI exception format. Some providers' mappers (e.g. + # _map_anthropic_exception, _map_aleph_alpha_exception) synchronously + # log a debug diagnostic (the raw status code) as part of mapping - + # restore the consumer's outer context only after this completes, so + # that diagnostic log line still carries the failing stream's own + # trace_id/session_id instead of the consumer's (or an empty one). if isinstance(e, OpenAIError): mapped_exception: Exception = e else: @@ -2152,6 +2251,7 @@ class CustomStreamWrapper: ) except Exception as mapping_error: mapped_exception = mapping_error + self._restore_consumer_correlation_context() def _normalize_status_code(exc: Exception) -> int | None: """Best-effort status_code extraction.""" diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index 36f3e875a7e..48d8a03d549 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,12 +1,13 @@ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from typing import ( TYPE_CHECKING, - Any, Final, TypeAlias, cast, ) +from typing_extensions import TypedDict + import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function @@ -39,6 +40,11 @@ _AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None" _ContextManagementSpec: TypeAlias = "dict[str, object] | list[dict[str, object]] | None" +class _CompletionKwargs(TypedDict, total=False, extra_items=object): + model: str + custom_llm_provider: str + + def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: @@ -312,7 +318,7 @@ ANTHROPIC_ADAPTER: Final = AnthropicAdapter() class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def _route_openai_thinking_to_responses_api_if_needed( - completion_kwargs: dict[str, Any], + completion_kwargs: _CompletionKwargs, *, thinking: Mapping[str, object] | None, ) -> None: @@ -377,7 +383,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def _normalize_reasoning_effort( - completion_kwargs: dict[str, Any], + completion_kwargs: _CompletionKwargs, ) -> None: """ Normalize reasoning_effort values based on target model capabilities. @@ -393,7 +399,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if reasoning_effort is None: return - model: Final = cast(str, completion_kwargs.get("model", "")) + model: Final = completion_kwargs.get("model", "") custom_llm_provider: Final = completion_kwargs.get("custom_llm_provider") if isinstance(reasoning_effort, str): @@ -417,19 +423,19 @@ class LiteLLMMessagesToCompletionTransformationHandler: max_tokens: int, messages: _AnthropicMessages, model: str, - metadata: dict | None = None, + metadata: dict[str, object] | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: _AnthropicSystem = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, + output_format: dict[str, object] | None = None, extra_kwargs: Mapping[str, object] | None = None, - ) -> tuple[dict[str, Any], dict[str, str]]: + ) -> tuple[_CompletionKwargs, dict[str, str]]: """Prepare kwargs for litellm.completion/acompletion. Returns: @@ -486,7 +492,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if openai_request is None: raise ValueError("Failed to translate request to OpenAI format") - completion_kwargs: Final[dict[str, Any]] = dict(openai_request) + completion_kwargs: Final[_CompletionKwargs] = {**openai_request} if stream: completion_kwargs["stream"] = stream @@ -538,17 +544,17 @@ class LiteLLMMessagesToCompletionTransformationHandler: max_tokens: int, messages: _AnthropicMessages, model: str, - metadata: dict | None = None, + metadata: dict[str, object] | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, + output_format: dict[str, object] | None = None, **kwargs, ) -> AnthropicMessagesResponse | AsyncIterator[bytes] | Iterator[bytes]: """Handle non-Anthropic models asynchronously using the adapter""" @@ -625,17 +631,17 @@ class LiteLLMMessagesToCompletionTransformationHandler: max_tokens: int, messages: _AnthropicMessages, model: str, - metadata: dict | None = None, + metadata: dict[str, object] | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, + output_format: dict[str, object] | None = None, _is_async: bool = False, **kwargs, ) -> ( 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 6ba129f5a7d..c1f10c245f8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/mcp_handler.py @@ -7,8 +7,8 @@ tool through a ``tool_use`` content block, and results are fed back as ``tool_result`` blocks in a user message. """ -from collections.abc import AsyncIterator, Mapping, Sequence -from typing import Any, Final +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence +from typing import Any, Final, NamedTuple from litellm._logging import verbose_logger from litellm.responses.mcp.request_context import MCPRequestContext @@ -24,14 +24,18 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( MAX_MCP_TOOL_USE_ITERATIONS: Final = 10 -def _get_response_content(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, Any]]: +class _AnthropicMessagesCall(NamedTuple): + fn: Callable[..., Awaitable[AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]]] + + +def _get_response_content(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, object]]: content: Final = response.get("content") if not isinstance(content, list): return () return tuple(block for block in content if isinstance(block, dict)) -def _extract_tool_use_blocks(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, Any]]: +def _extract_tool_use_blocks(response: AnthropicMessagesResponse) -> Sequence[Mapping[str, object]]: """Return the ``tool_use`` content blocks the model emitted.""" return tuple(block for block in _get_response_content(response) if block.get("type") == "tool_use") @@ -41,7 +45,7 @@ def _get_stop_reason(response: AnthropicMessagesResponse) -> str | None: return stop_reason if isinstance(stop_reason, str) else None -def _build_tool_result_message(tool_results: Sequence[Mapping[str, Any]]) -> AnthropicMessagesUserMessageParam: +def _build_tool_result_message(tool_results: Sequence[Mapping[str, object]]) -> AnthropicMessagesUserMessageParam: """Turn executed tool results into the user message Anthropic expects.""" return AnthropicMessagesUserMessageParam( role="user", @@ -58,11 +62,11 @@ def _build_tool_result_message(tool_results: Sequence[Mapping[str, Any]]) -> Ant async def anthropic_messages_with_mcp( max_tokens: int, - messages: Sequence[Mapping[str, Any]], + messages: Sequence[Mapping[str, object]], model: str, - tools: Sequence[Mapping[str, Any]] | None = None, + tools: Sequence[Mapping[str, object]] | None = None, **kwargs: Any, # kwargs-ok: forwarded verbatim to litellm.anthropic_messages, which owns the param contract -) -> AnthropicMessagesResponse | AsyncIterator[Any]: +) -> AnthropicMessagesResponse | Iterator[bytes] | AsyncIterator[object]: """ Expand litellm_proxy MCP references for `/v1/messages` and run the tool loop. @@ -81,7 +85,7 @@ async def anthropic_messages_with_mcp( mcp_references, other_tools = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) if not mcp_references: - return await litellm.anthropic_messages( + return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( max_tokens=max_tokens, messages=list(messages), model=model, @@ -114,7 +118,7 @@ async def anthropic_messages_with_mcp( ) stream: Final = bool(kwargs.pop("stream", False)) - base_call_args: Final[Mapping[str, Any]] = { + base_call_args: Final[Mapping[str, object]] = { "max_tokens": max_tokens, "model": model, "tools": all_tools or None, @@ -123,10 +127,12 @@ async def anthropic_messages_with_mcp( } if not should_auto_execute: - return await litellm.anthropic_messages(messages=list(messages), stream=stream, **base_call_args) + return await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( + messages=list(messages), stream=stream, **base_call_args + ) - working_messages: Sequence[Mapping[str, Any]] = tuple(messages) - response: AnthropicMessagesResponse = await litellm.anthropic_messages( + working_messages: Sequence[Mapping[str, object]] = tuple(messages) + response: AnthropicMessagesResponse = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( messages=list(working_messages), stream=False, **base_call_args ) @@ -161,7 +167,9 @@ async def anthropic_messages_with_mcp( {"role": "assistant", "content": list(_get_response_content(response))}, _build_tool_result_message(tool_results), ) - response = await litellm.anthropic_messages(messages=list(working_messages), stream=False, **base_call_args) + response = await _AnthropicMessagesCall(fn=litellm.anthropic_messages).fn( + messages=list(working_messages), stream=False, **base_call_args + ) else: verbose_logger.warning( "MCP tool loop hit its %s iteration cap for model %s; returning the last response", diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index 5e05ebc3c63..9210719dd59 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -8,7 +8,12 @@ from collections.abc import AsyncIterator, Coroutine from typing import Any, Final import litellm -from litellm.types.llms.anthropic import AnthropicMessagesRequest +from litellm.types.llms.anthropic import ( + AllAnthropicToolsValues, + AnthropicMessagesRequest, + AnthropicOutputConfig, + AnthropicOutputSchema, +) from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -27,24 +32,24 @@ def _build_responses_kwargs( model: str, context_management: dict | None = None, metadata: dict | None = None, - output_config: dict | None = None, + output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[AllAnthropicToolsValues | dict] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, + output_format: AnthropicOutputSchema | None = None, extra_kwargs: dict[str, Any] | None = None, ) -> dict[str, Any]: """ Build the kwargs dict to pass directly to litellm.responses() / litellm.aresponses(). """ # Build a typed AnthropicMessagesRequest for the adapter - request_data: Final[dict[str, Any]] = { + request_data: Final[AnthropicMessagesRequest] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -128,19 +133,19 @@ class LiteLLMMessagesToResponsesAPIHandler: model: str, context_management: dict | None = None, metadata: dict | None = None, - output_config: dict | None = None, + output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[AllAnthropicToolsValues | dict] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, + output_format: AnthropicOutputSchema | None = None, **kwargs, - ) -> AnthropicMessagesResponse | AsyncIterator: + ) -> AnthropicMessagesResponse | AsyncIterator[bytes]: responses_kwargs: Final = _build_responses_kwargs( max_tokens=max_tokens, messages=messages, @@ -179,23 +184,23 @@ class LiteLLMMessagesToResponsesAPIHandler: model: str, context_management: dict | None = None, metadata: dict | None = None, - output_config: dict | None = None, + output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[AllAnthropicToolsValues | dict] | None = None, top_k: int | None = None, top_p: float | None = None, - output_format: dict | None = None, + output_format: AnthropicOutputSchema | None = None, _is_async: bool = False, **kwargs, ) -> ( AnthropicMessagesResponse - | AsyncIterator[Any] - | Coroutine[Any, Any, AnthropicMessagesResponse | AsyncIterator[Any]] + | AsyncIterator[bytes] + | Coroutine[None, None, AnthropicMessagesResponse | AsyncIterator[bytes]] ): if _is_async: return LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( 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 85fda3a6522..fad7e7558c2 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -67,6 +67,8 @@ class AmazonAnthropicClaudeMessagesConfig( DEFAULT_BEDROCK_ANTHROPIC_API_VERSION = "bedrock-2023-05-31" + WEBSEARCH_INTERCEPTION_DOCS_URL = "https://docs.litellm.ai/docs/integrations/websearch_interception" + @property def custom_llm_provider(self) -> str | None: return "bedrock" @@ -572,6 +574,45 @@ class AmazonAnthropicClaudeMessagesConfig( return filtered_betas + @staticmethod + def _reject_unsupported_web_search_tools(anthropic_messages_request: dict[str, object], model: str) -> None: + """ + Bedrock's Anthropic endpoints cannot execute Anthropic's server-side + ``web_search_*`` tool; forwarding it returns an opaque + "The provided request is not valid" 400 from Bedrock. Fail fast with an + error that names the problem and the fix instead. + + When web search interception is enabled + (``litellm_settings.callbacks: ["websearch_interception"]``), the tool + is converted to a regular function tool before this transform runs, so + this guard never fires. + """ + from litellm.integrations.websearch_interception.tools import ( + is_anthropic_native_web_search_tool, + ) + + tools: Final = anthropic_messages_request.get("tools") + if not isinstance(tools, list): + return + web_search_tool: Final = next( + (t for t in tools if isinstance(t, dict) and is_anthropic_native_web_search_tool(t)), + None, + ) + if web_search_tool is None: + return + raise litellm.BadRequestError( + message=( + f"Bedrock does not support Anthropic's server-side web search tool " + f"(tool type '{web_search_tool.get('type')}', model '{model}'). " + "To use web search with this model, enable LiteLLM's web search interception " + "so the proxy executes the search instead: " + f"{AmazonAnthropicClaudeMessagesConfig.WEBSEARCH_INTERCEPTION_DOCS_URL}. " + "Alternatively, remove the web_search tool from the request." + ), + model=model, + llm_provider="bedrock", + ) + def _strip_unsupported_bedrock_invoke_fields( self, anthropic_messages_request: dict, @@ -630,6 +671,8 @@ class AmazonAnthropicClaudeMessagesConfig( ############## BEDROCK Invoke SPECIFIC TRANSFORMATION ### ######################################################### + self._reject_unsupported_web_search_tools(anthropic_messages_request=anthropic_messages_request, model=model) + # 1. anthropic_version is required for all claude models if "anthropic_version" not in anthropic_messages_request: anthropic_messages_request["anthropic_version"] = self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d954e33da9c..951c114b0a9 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -2808,6 +2808,7 @@ }, "azure/codex-mini": { "cache_read_input_token_cost": 3.75e-07, + "deprecation_date": "2026-11-15", "input_cost_per_token": 1.5e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -3627,7 +3628,7 @@ "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, "azure/eu/gpt-4o-2024-08-06": { - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.375e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -3644,7 +3645,7 @@ "supports_vision": true }, "azure/eu/gpt-4o-2024-11-20": { - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "cache_creation_input_token_cost": 1.38e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -3661,6 +3662,7 @@ }, "azure/eu/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 8.3e-08, + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -3742,6 +3744,7 @@ }, "azure/eu/gpt-5-2025-08-07": { "cache_read_input_token_cost": 1.375e-07, + "deprecation_date": "2027-02-09", "input_cost_per_token": 1.375e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -3774,6 +3777,7 @@ }, "azure/eu/gpt-5-mini-2025-08-07": { "cache_read_input_token_cost": 2.75e-08, + "deprecation_date": "2027-02-09", "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -3840,6 +3844,7 @@ }, "azure/eu/gpt-5.1-chat": { "cache_read_input_token_cost": 1.4e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -3934,6 +3939,7 @@ }, "azure/eu/gpt-5-nano-2025-08-07": { "cache_read_input_token_cost": 5.5e-09, + "deprecation_date": "2027-02-09", "input_cost_per_token": 5.5e-08, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -3966,6 +3972,7 @@ }, "azure/eu/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -4011,6 +4018,7 @@ }, "azure/eu/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, + "deprecation_date": "2026-10-01", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -4027,7 +4035,7 @@ }, "azure/global-standard/gpt-4o-2024-08-06": { "cache_read_input_token_cost": 1.25e-06, - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4044,7 +4052,7 @@ }, "azure/global-standard/gpt-4o-2024-11-20": { "cache_read_input_token_cost": 1.25e-06, - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4073,7 +4081,7 @@ "supports_vision": true }, "azure/global/gpt-4o-2024-08-06": { - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4090,7 +4098,7 @@ "supports_vision": true }, "azure/global/gpt-4o-2024-11-20": { - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4142,6 +4150,7 @@ }, "azure/global/gpt-5.1-chat": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4476,7 +4485,7 @@ "supports_web_search": false }, "azure/gpt-4.1-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -4543,7 +4552,7 @@ "supports_web_search": false }, "azure/gpt-4.1-mini-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -4609,7 +4618,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -4677,6 +4686,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-05-13": { + "deprecation_date": "2026-10-01", "input_cost_per_token": 5e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4691,7 +4701,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-08-06": { - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4708,7 +4718,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-11-20": { - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -4725,6 +4735,7 @@ "supports_vision": true }, "azure/gpt-audio-2025-08-28": { + "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4756,6 +4767,7 @@ "supports_vision": false }, "azure/gpt-audio-1.5-2026-02-23": { + "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4787,6 +4799,7 @@ "supports_vision": false }, "azure/gpt-audio-mini-2025-10-06": { + "deprecation_date": "2027-04-06", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "azure", @@ -4866,6 +4879,7 @@ }, "azure/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4933,6 +4947,7 @@ "azure/gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_token_cost": 4e-06, + "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, @@ -4965,6 +4980,7 @@ "azure/gpt-realtime-1.5-2026-02-23": { "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_token_cost": 4e-06, + "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, @@ -5102,6 +5118,7 @@ "supports_tool_choice": true }, "azure/gpt-4o-transcribe": { + "deprecation_date": "2026-10-15", "input_cost_per_audio_token": 2.5e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -5114,6 +5131,7 @@ ] }, "azure/gpt-4o-transcribe-diarize": { + "deprecation_date": "2027-04-15", "input_cost_per_audio_token": 2.5e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -5145,6 +5163,7 @@ "azure/gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2027-05-15", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "azure", @@ -5182,6 +5201,7 @@ "azure/gpt-5.1-chat-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "azure", @@ -5218,6 +5238,7 @@ "azure/gpt-5.1-codex-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2027-05-15", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "azure", @@ -5251,6 +5272,7 @@ "azure/gpt-5.1-codex-mini-2025-11-13": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2027-05-15", "input_cost_per_token": 2.5e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "azure", @@ -5315,6 +5337,7 @@ }, "azure/gpt-5-2025-08-07": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2027-02-09", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5347,6 +5370,7 @@ }, "azure/gpt-5-chat": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -5380,6 +5404,7 @@ }, "azure/gpt-5-chat-latest": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -5412,6 +5437,7 @@ }, "azure/gpt-5-codex": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2027-03-17", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5474,6 +5500,7 @@ }, "azure/gpt-5-mini-2025-08-07": { "cache_read_input_token_cost": 2.5e-08, + "deprecation_date": "2027-02-09", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5538,6 +5565,7 @@ }, "azure/gpt-5-nano-2025-08-07": { "cache_read_input_token_cost": 5e-09, + "deprecation_date": "2027-02-09", "input_cost_per_token": 5e-08, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5569,6 +5597,7 @@ "supports_vision": true }, "azure/gpt-5-pro": { + "deprecation_date": "2027-04-07", "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5633,6 +5662,7 @@ }, "azure/gpt-5.1-chat": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -5697,6 +5727,7 @@ }, "azure/gpt-5.1-codex-max": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2027-05-18", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5791,6 +5822,7 @@ "azure/gpt-5.2-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2027-06-08", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "azure", @@ -5827,6 +5859,7 @@ "azure/gpt-5.2-chat": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "azure", @@ -5861,6 +5894,7 @@ "azure/gpt-5.2-chat-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-05-13", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "azure", @@ -5894,6 +5928,7 @@ }, "azure/gpt-5.2-codex": { "cache_read_input_token_cost": 1.75e-07, + "deprecation_date": "2027-07-13", "input_cost_per_token": 1.75e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5925,6 +5960,7 @@ "azure/gpt-5.3-chat": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "azure", @@ -5958,6 +5994,7 @@ }, "azure/gpt-5.3-codex": { "cache_read_input_token_cost": 1.75e-07, + "deprecation_date": "2027-08-24", "input_cost_per_token": 1.75e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -6164,6 +6201,7 @@ "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, "cache_read_input_token_cost_above_272k_tokens_priority": 1e-06, + "deprecation_date": "2027-09-02", "input_cost_per_token": 2.5e-06, "input_cost_per_token_above_272k_tokens": 5e-06, "input_cost_per_token_priority": 5e-06, @@ -6203,6 +6241,7 @@ "azure/us/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, + "deprecation_date": "2027-09-02", "input_cost_per_token": 2.75e-06, "input_cost_per_token_priority": 5.5e-06, "output_cost_per_token": 1.65e-05, @@ -6238,6 +6277,7 @@ "azure/eu/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, + "deprecation_date": "2027-09-02", "input_cost_per_token": 2.75e-06, "input_cost_per_token_priority": 5.5e-06, "output_cost_per_token": 1.65e-05, @@ -6308,6 +6348,7 @@ "azure/gpt-5.4-pro-2026-03-05": { "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "deprecation_date": "2027-09-07", "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", @@ -6390,6 +6431,7 @@ "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, "input_cost_per_token_priority": 1e-05, @@ -6435,6 +6477,7 @@ "cache_read_input_token_cost_above_272k_tokens": 4e-07, "cache_read_input_token_cost_priority": 4e-07, "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "input_cost_per_token_priority": 4e-06, @@ -6480,6 +6523,7 @@ "cache_read_input_token_cost_above_272k_tokens": 4e-08, "cache_read_input_token_cost_priority": 4e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, "input_cost_per_token_priority": 4e-07, @@ -6566,6 +6610,7 @@ "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, @@ -6608,6 +6653,7 @@ "cache_read_input_token_cost": 2.2e-07, "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "cache_read_input_token_cost_priority": 5.5e-07, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2.2e-06, "input_cost_per_token_above_272k_tokens": 4.4e-06, "input_cost_per_token_priority": 5.5e-06, @@ -6650,6 +6696,7 @@ "cache_read_input_token_cost": 2.2e-08, "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, "input_cost_per_token_priority": 5.5e-07, @@ -6734,6 +6781,7 @@ "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, @@ -6776,6 +6824,7 @@ "cache_read_input_token_cost": 2.2e-07, "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "cache_read_input_token_cost_priority": 5.5e-07, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2.2e-06, "input_cost_per_token_above_272k_tokens": 4.4e-06, "input_cost_per_token_priority": 5.5e-06, @@ -6818,6 +6867,7 @@ "cache_read_input_token_cost": 2.2e-08, "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, "input_cost_per_token_priority": 5.5e-07, @@ -7216,6 +7266,7 @@ }, "azure/gpt-5.4-mini-2026-03-17": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2027-09-21", "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -7286,6 +7337,7 @@ }, "azure/gpt-5.4-nano-2026-03-17": { "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2027-09-21", "input_cost_per_token": 2e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -7321,6 +7373,7 @@ }, "azure/gpt-image-1": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-10-23", "input_cost_per_image_token": 1e-05, "input_cost_per_token": 5e-06, "litellm_provider": "azure", @@ -7432,6 +7485,7 @@ }, "azure/gpt-image-1-mini": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2027-04-07", "input_cost_per_image_token": 2.5e-06, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -7456,6 +7510,7 @@ }, "azure/gpt-image-1.5-2025-12-16": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-06-16", "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, "litellm_provider": "azure", @@ -7483,6 +7538,7 @@ }, "azure/gpt-image-2-2026-04-21": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-10-21", "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, "litellm_provider": "azure", @@ -7613,6 +7669,7 @@ }, "azure/o1-2024-12-17": { "cache_read_input_token_cost": 7.5e-06, + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -7718,7 +7775,7 @@ "supports_vision": true }, "azure/o3-2025-04-16": { - "deprecation_date": "2026-04-16", + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -7749,6 +7806,7 @@ }, "azure/o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2026-12-26", "input_cost_per_token": 1e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -7796,6 +7854,7 @@ }, "azure/o3-mini-2025-01-31": { "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-10-01", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -7839,6 +7898,7 @@ "supports_vision": true }, "azure/o3-pro-2025-06-10": { + "deprecation_date": "2026-12-17", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -7899,6 +7959,7 @@ }, "azure/o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, + "deprecation_date": "2026-10-16", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -7939,6 +8000,7 @@ "output_cost_per_token": 0.0 }, "azure/text-embedding-3-large": { + "deprecation_date": "2028-02-09", "input_cost_per_token": 1.3e-07, "litellm_provider": "azure", "max_input_tokens": 8191, @@ -7947,7 +8009,7 @@ "output_cost_per_token": 0.0 }, "azure/text-embedding-3-small": { - "deprecation_date": "2026-04-30", + "deprecation_date": "2028-02-09", "input_cost_per_token": 2e-08, "litellm_provider": "azure", "max_input_tokens": 8191, @@ -7956,6 +8018,7 @@ "output_cost_per_token": 0.0 }, "azure/text-embedding-ada-002": { + "deprecation_date": "2028-02-09", "input_cost_per_token": 1e-07, "litellm_provider": "azure", "max_input_tokens": 8191, @@ -7987,17 +8050,19 @@ ] }, "azure/tts-1": { + "deprecation_date": "2026-12-15", "input_cost_per_character": 1.5e-05, "litellm_provider": "azure", "mode": "audio_speech" }, "azure/tts-1-hd": { + "deprecation_date": "2026-12-15", "input_cost_per_character": 3e-05, "litellm_provider": "azure", "mode": "audio_speech" }, "azure/us/gpt-4.1-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "input_cost_per_token_batches": 1.1e-06, @@ -8031,7 +8096,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-mini-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 4.4e-07, "input_cost_per_token_batches": 2.2e-07, @@ -8065,7 +8130,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 6e-08, @@ -8098,7 +8163,7 @@ "supports_vision": true }, "azure/us/gpt-4o-2024-08-06": { - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.375e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -8115,7 +8180,7 @@ "supports_vision": true }, "azure/us/gpt-4o-2024-11-20": { - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "cache_creation_input_token_cost": 1.38e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -8132,6 +8197,7 @@ }, "azure/us/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 8.3e-08, + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -8213,6 +8279,7 @@ }, "azure/us/gpt-5-2025-08-07": { "cache_read_input_token_cost": 1.375e-07, + "deprecation_date": "2027-02-09", "input_cost_per_token": 1.375e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -8245,6 +8312,7 @@ }, "azure/us/gpt-5-mini-2025-08-07": { "cache_read_input_token_cost": 2.75e-08, + "deprecation_date": "2027-02-09", "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -8277,6 +8345,7 @@ }, "azure/us/gpt-5-nano-2025-08-07": { "cache_read_input_token_cost": 5.5e-09, + "deprecation_date": "2027-02-09", "input_cost_per_token": 5.5e-08, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -8343,6 +8412,7 @@ }, "azure/us/gpt-5.1-chat": { "cache_read_input_token_cost": 1.4e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -8437,6 +8507,7 @@ }, "azure/us/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8481,7 +8552,7 @@ "supports_vision": false }, "azure/us/o3-2025-04-16": { - "deprecation_date": "2026-04-16", + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "azure", @@ -8512,6 +8583,7 @@ }, "azure/us/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, + "deprecation_date": "2026-10-01", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -8528,6 +8600,7 @@ }, "azure/us/o4-mini-2025-04-16": { "cache_read_input_token_cost": 3.1e-07, + "deprecation_date": "2026-10-16", "input_cost_per_token": 1.21e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8544,6 +8617,7 @@ "supports_vision": true }, "azure/whisper-1": { + "deprecation_date": "2026-12-15", "input_cost_per_second": 0.0001, "litellm_provider": "azure", "mode": "audio_transcription", @@ -11396,6 +11470,7 @@ "output_cost_per_token": 5e-07 }, "chatgpt-4o-latest": { + "deprecation_date": "2026-02-17", "input_cost_per_token": 5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -11499,6 +11574,7 @@ "cache_creation_input_token_cost": 3e-07, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-04-20", "input_cost_per_token": 2.5e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, @@ -11517,7 +11593,7 @@ "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 1.5e-06, - "deprecation_date": "2026-05-01", + "deprecation_date": "2026-01-05", "input_cost_per_token": 1.5e-05, "litellm_provider": "anthropic", "max_input_tokens": 200000, @@ -11535,6 +11611,7 @@ "claude-4-opus-20250514": { "cache_creation_input_token_cost": 1.875e-05, "cache_read_input_token_cost": 1.5e-06, + "deprecation_date": "2026-06-15", "input_cost_per_token": 1.5e-05, "litellm_provider": "anthropic", "max_input_tokens": 200000, @@ -11563,6 +11640,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "deprecation_date": "2026-06-15", "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "litellm_provider": "anthropic", @@ -11776,7 +11854,8 @@ "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "deprecation_date": "2026-08-05" }, "claude-opus-4-1-20250805": { "cache_creation_input_token_cost": 1.875e-05, @@ -11812,7 +11891,7 @@ "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, "input_cost_per_token": 1.5e-05, - "deprecation_date": "2026-05-14", + "deprecation_date": "2026-06-15", "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -12153,7 +12232,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-20250514": { - "deprecation_date": "2026-05-14", + "deprecation_date": "2026-06-15", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -12508,6 +12587,7 @@ }, "codex-mini-latest": { "cache_read_input_token_cost": 3.75e-07, + "deprecation_date": "2026-02-12", "input_cost_per_token": 1.5e-06, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -12746,6 +12826,7 @@ "supports_vision": true }, "dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_image": 0.02, "litellm_provider": "openai", "mode": "image_generation", @@ -12756,6 +12837,7 @@ ] }, "dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_image": 0.04, "litellm_provider": "openai", "mode": "image_generation", @@ -15471,6 +15553,7 @@ "output_cost_per_token": 1.85e-06, "supports_function_calling": true, "supports_reasoning": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, @@ -17259,6 +17342,7 @@ "supports_tool_choice": true }, "ft:babbage-002": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, "input_cost_per_token_batches": 2e-07, "litellm_provider": "text-completion-openai", @@ -17270,6 +17354,7 @@ "output_cost_per_token_batches": 2e-07 }, "ft:davinci-002": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.2e-05, "input_cost_per_token_batches": 1e-06, "litellm_provider": "text-completion-openai", @@ -17281,6 +17366,7 @@ "output_cost_per_token_batches": 1e-06 }, "ft:gpt-3.5-turbo": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "openai", @@ -17294,6 +17380,7 @@ "supports_tool_choice": true }, "ft:gpt-3.5-turbo-0125": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -17327,6 +17414,7 @@ "supports_tool_choice": true }, "ft:gpt-4-0613": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-05, "litellm_provider": "openai", "max_input_tokens": 8192, @@ -17433,6 +17521,7 @@ }, "ft:gpt-4.1-nano-2025-04-14": { "cache_read_input_token_cost": 5e-08, + "deprecation_date": "2026-10-23", "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", @@ -17451,6 +17540,7 @@ }, "ft:o4-mini-2025-04-16": { "cache_read_input_token_cost": 1e-06, + "deprecation_date": "2026-10-23", "input_cost_per_token": 4e-06, "input_cost_per_token_batches": 2e-06, "litellm_provider": "openai", @@ -17702,6 +17792,46 @@ "tpm": 8000000, "supports_image_size": false }, + "gemini-3-pro-image": { + "input_cost_per_image": 0.0011, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.134, + "output_cost_per_image_token": 0.00012, + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3-pro-image", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, @@ -17742,6 +17872,44 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-3.1-flash-image": { + "input_cost_per_image": 0.00056, + "input_cost_per_token": 5e-07, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.0672, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 3e-06, + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "web_search_billing_unit": "per_query" + }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, "input_cost_per_token": 5e-07, @@ -18847,6 +19015,7 @@ }, "gemini/gemini-robotics-er-1.5-preview": { "cache_read_input_token_cost": 0, + "deprecation_date": "2026-04-30", "input_cost_per_token": 3e-07, "input_cost_per_audio_token": 1e-06, "litellm_provider": "gemini", @@ -19089,6 +19258,7 @@ "uses_embed_content": true }, "gemini/gemini-embedding-001": { + "deprecation_date": "2028-05-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "gemini", "max_input_tokens": 2048, @@ -19101,6 +19271,7 @@ "tpm": 10000000 }, "gemini/gemini-embedding-2-preview": { + "deprecation_date": "2026-08-10", "input_cost_per_audio_per_second": 0.00016, "input_cost_per_image": 0.00012, "input_cost_per_token": 2e-07, @@ -19310,6 +19481,7 @@ }, "gemini/gemini-2.5-flash-image": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-02", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -19389,7 +19561,6 @@ ], "supports_function_calling": false, "supports_prompt_caching": true, - "supports_reasoning": false, "supports_response_schema": true, "supports_system_messages": true, "supports_vision": true, @@ -19399,9 +19570,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "supports_reasoning": false }, "gemini/gemini-3-pro-image-preview": { + "deprecation_date": "2026-06-25", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -19487,6 +19660,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-flash-image-preview": { + "deprecation_date": "2026-06-25", "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "litellm_provider": "gemini", @@ -19618,6 +19792,7 @@ }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { "cache_read_input_token_cost": 1e-08, + "deprecation_date": "2026-03-31", "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -19999,6 +20174,7 @@ }, "gemini/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, + "deprecation_date": "2026-05-25", "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "gemini", @@ -20051,6 +20227,7 @@ "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2027-05-07", "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -20818,18 +20995,21 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "gemini/imagen-4.0-fast-generate-001": { + "deprecation_date": "2026-08-17", "litellm_provider": "gemini", "mode": "image_generation", "output_cost_per_image": 0.02, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "gemini/imagen-4.0-generate-001": { + "deprecation_date": "2026-08-17", "litellm_provider": "gemini", "mode": "image_generation", "output_cost_per_image": 0.04, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "gemini/imagen-4.0-ultra-generate-001": { + "deprecation_date": "2026-08-17", "litellm_provider": "gemini", "mode": "image_generation", "output_cost_per_image": 0.06, @@ -20913,6 +21093,7 @@ "supports_web_search": false }, "gemini/veo-2.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -21004,8 +21185,7 @@ "max_tokens": 16000, "mode": "chat", "supported_endpoints": [ - "/v1/chat/completions", - "/v1/messages" + "/v1/chat/completions" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -21018,8 +21198,7 @@ "max_tokens": 16000, "mode": "chat", "supported_endpoints": [ - "/v1/chat/completions", - "/v1/messages" + "/v1/chat/completions" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -21071,8 +21250,7 @@ "max_tokens": 16000, "mode": "chat", "supported_endpoints": [ - "/v1/chat/completions", - "/v1/messages" + "/v1/chat/completions" ], "supports_function_calling": true, "supports_parallel_function_calling": true, @@ -21950,6 +22128,7 @@ "supports_vision": true }, "gpt-3.5-turbo": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 5e-07, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -21963,6 +22142,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-0125": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 5e-07, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -22022,6 +22202,7 @@ "output_cost_per_token": 2e-06 }, "gpt-4": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-05, "litellm_provider": "openai", "max_input_tokens": 8192, @@ -22062,7 +22243,7 @@ "supports_tool_choice": true }, "gpt-4-0613": { - "deprecation_date": "2025-06-06", + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-05, "litellm_provider": "openai", "max_input_tokens": 8192, @@ -22076,7 +22257,7 @@ "supports_tool_choice": true }, "gpt-4-1106-preview": { - "deprecation_date": "2026-03-26", + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22091,6 +22272,7 @@ "supports_tool_choice": true }, "gpt-4-turbo": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22107,6 +22289,7 @@ "supports_vision": true }, "gpt-4-turbo-2024-04-09": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22288,6 +22471,7 @@ "gpt-4.1-nano": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 5e-08, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, "input_cost_per_token_priority": 2e-07, @@ -22324,6 +22508,7 @@ "gpt-4.1-nano-2025-04-14": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 5e-08, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-07, "input_cost_per_token_priority": 2e-07, "input_cost_per_token_batches": 5e-08, @@ -22381,6 +22566,7 @@ "supports_vision": true }, "gpt-4o-2024-05-13": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_priority": 8.75e-06, @@ -22447,6 +22633,7 @@ "supports_vision": true }, "gpt-4o-audio-preview": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22464,6 +22651,7 @@ "supports_tool_choice": true }, "gpt-4o-audio-preview-2024-12-17": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22481,6 +22669,7 @@ "supports_tool_choice": true }, "gpt-4o-audio-preview-2025-06-03": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22498,6 +22687,7 @@ "supports_tool_choice": true }, "gpt-audio": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22567,6 +22757,7 @@ "supports_vision": false }, "gpt-audio-2025-08-28": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22603,6 +22794,7 @@ "supports_vision": false }, "gpt-audio-mini": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -22639,6 +22831,7 @@ "supports_vision": false }, "gpt-audio-mini-2025-10-06": { + "deprecation_date": "2026-07-23", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -22762,6 +22955,7 @@ "supports_vision": true }, "gpt-4o-mini-audio-preview": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 1.5e-07, "litellm_provider": "openai", @@ -22779,6 +22973,7 @@ "supports_tool_choice": true }, "gpt-4o-mini-audio-preview-2024-12-17": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 1.5e-07, "litellm_provider": "openai", @@ -22798,6 +22993,7 @@ "gpt-4o-mini-realtime-preview": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 3e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -22817,6 +23013,7 @@ "gpt-4o-mini-realtime-preview-2024-12-17": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 3e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -22861,6 +23058,7 @@ }, "gpt-4o-mini-search-preview-2025-03-11": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "openai", @@ -22911,6 +23109,7 @@ }, "gpt-4o-realtime-preview": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -22929,6 +23128,7 @@ }, "gpt-4o-realtime-preview-2024-12-17": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -22947,6 +23147,7 @@ }, "gpt-4o-realtime-preview-2025-06-03": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -22991,6 +23192,7 @@ }, "gpt-4o-search-preview-2025-03-11": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-07-23", "input_cost_per_token": 2.5e-06, "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "openai", @@ -23023,6 +23225,7 @@ }, "gpt-image-1.5": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-12-01", "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", @@ -23037,6 +23240,7 @@ }, "gpt-image-1.5-2025-12-16": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-12-01", "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", @@ -23532,6 +23736,7 @@ "gpt-5.1-chat-latest": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -23651,6 +23856,7 @@ "gpt-5.2-chat-latest": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-08-10", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -23689,6 +23895,7 @@ "gpt-5.3-chat-latest": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-08-10", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -24571,7 +24778,7 @@ "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, + "max_input_tokens": 400000, "max_output_tokens": 272000, "max_tokens": 272000, "mode": "responses", @@ -24604,10 +24811,11 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-pro-2025-10-06": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", - "max_input_tokens": 128000, + "max_input_tokens": 400000, "max_output_tokens": 272000, "max_tokens": 272000, "mode": "responses", @@ -24643,6 +24851,7 @@ "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.25e-06, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, @@ -24718,6 +24927,7 @@ }, "gpt-5-chat-latest": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -24753,6 +24963,7 @@ }, "gpt-5-codex": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -24788,6 +24999,7 @@ "gpt-5.1-codex": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -24824,6 +25036,7 @@ }, "gpt-5.1-codex-max": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -24859,6 +25072,7 @@ "gpt-5.1-codex-mini": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-07-23", "input_cost_per_token": 2.5e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "openai", @@ -24896,6 +25110,7 @@ "gpt-5.2-codex": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -25013,6 +25228,7 @@ "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2.5e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, @@ -25094,6 +25310,7 @@ "gpt-5-nano-2025-08-07": { "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, + "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, "input_cost_per_token_priority": 2.5e-06, "input_cost_per_token_flex": 2.5e-08, @@ -25133,6 +25350,7 @@ }, "gpt-image-1": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-10-23", "input_cost_per_image_token": 1e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -25145,6 +25363,7 @@ }, "gpt-image-1-mini": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2026-12-01", "input_cost_per_image_token": 2.5e-06, "input_cost_per_token": 2e-06, "litellm_provider": "openai", @@ -25158,6 +25377,7 @@ "gpt-realtime": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, @@ -25324,6 +25544,7 @@ "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -25355,6 +25576,7 @@ "gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, @@ -29183,6 +29405,7 @@ }, "o1": { "cache_read_input_token_cost": 7.5e-06, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -29202,6 +29425,7 @@ }, "o1-2024-12-17": { "cache_read_input_token_cost": 7.5e-06, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -29220,6 +29444,7 @@ "supports_vision": true }, "o1-pro": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 0.00015, "input_cost_per_token_batches": 7.5e-05, "litellm_provider": "openai", @@ -29252,6 +29477,7 @@ "supports_vision": true }, "o1-pro-2025-03-19": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 0.00015, "input_cost_per_token_batches": 7.5e-05, "litellm_provider": "openai", @@ -29325,6 +29551,7 @@ "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 3.5e-06, @@ -29361,6 +29588,7 @@ }, "o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, "litellm_provider": "openai", @@ -29395,6 +29623,7 @@ }, "o3-deep-research-2025-06-26": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, "litellm_provider": "openai", @@ -29429,6 +29658,7 @@ }, "o3-mini": { "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -29446,6 +29676,7 @@ }, "o3-mini-2025-01-31": { "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -29493,6 +29724,7 @@ "supports_web_search": true }, "o3-pro-2025-06-10": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -29527,6 +29759,7 @@ "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.375e-07, "cache_read_input_token_cost_priority": 5e-07, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, "input_cost_per_token_flex": 5.5e-07, "input_cost_per_token_priority": 2e-06, @@ -29552,6 +29785,7 @@ "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.375e-07, "cache_read_input_token_cost_priority": 5e-07, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, "input_cost_per_token_flex": 5.5e-07, "input_cost_per_token_priority": 2e-06, @@ -29575,6 +29809,7 @@ }, "o4-mini-deep-research": { "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", @@ -29609,6 +29844,7 @@ }, "o4-mini-deep-research-2025-06-26": { "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", @@ -32000,6 +32236,22 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/z-ai/glm-5.1": { + "input_cost_per_token": 1.05e-06, + "output_cost_per_token": 3.5e-06, + "cache_read_input_token_cost": 5.25e-07, + "cache_creation_input_token_cost": 0.0, + "litellm_provider": "openrouter", + "max_input_tokens": 202752, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "source": "https://openrouter.ai/z-ai/glm-5.1", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/minimax/minimax-m2.1": { "input_cost_per_token": 2.7e-07, "output_cost_per_token": 1.2e-06, @@ -39994,69 +40246,86 @@ }, "xai/grok-4.20-multi-agent-beta-0309": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true }, "xai/grok-4.20-beta-0309-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true }, "xai/grok-4.20-0309-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_prompt_caching": true, + "supports_response_schema": true }, "xai/grok-4.20-beta-0309-non-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true }, "xai/grok-4.3": { "cache_read_input_token_cost": 2e-07, @@ -40101,8 +40370,8 @@ "supports_web_search": true }, "xai/grok-4.5": { - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -40122,8 +40391,8 @@ "supports_web_search": true }, "xai/grok-4.5-latest": { - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -40156,51 +40425,64 @@ "supports_web_search": true }, "xai/grok-code-fast": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 2e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 2e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true, + "supports_vision": true }, "xai/grok-code-fast-1": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 2e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 2e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "deprecation_date": "2026-05-15" + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true, + "supports_vision": true }, "xai/grok-code-fast-1-0825": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 2e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 2e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "deprecation_date": "2026-05-15" + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true, + "supports_vision": true }, "xai/grok-vision-beta": { "input_cost_per_image": 5e-06, @@ -40240,6 +40522,7 @@ "mode": "chat", "supports_function_calling": true, "supports_reasoning": true, + "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" @@ -40273,6 +40556,21 @@ "supports_tool_choice": true, "source": "https://docs.z.ai/guides/overview/pricing" }, + "zai/glm-5.1": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 2.6e-07, + "input_cost_per_token": 1.4e-06, + "output_cost_per_token": 4.4e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-5-code": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 3e-07, @@ -40303,6 +40601,21 @@ "supports_tool_choice": true, "source": "https://docs.z.ai/guides/overview/pricing" }, + "zai/glm-4.7-flash": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 0, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.6": { "cache_creation_input_token_cost": 0, "cache_read_input_token_cost": 1.1e-07, @@ -40407,6 +40720,7 @@ "mode": "chat" }, "openai/sora-2": { + "deprecation_date": "2026-09-24", "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.1, @@ -40420,6 +40734,7 @@ ] }, "openai/sora-2-pro": { + "deprecation_date": "2026-09-24", "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.3, @@ -44338,6 +44653,7 @@ ] }, "gpt-4o-mini-tts-2025-03-20": { + "deprecation_date": "2026-07-23", "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "mode": "audio_speech", @@ -44374,6 +44690,7 @@ ] }, "gpt-4o-mini-transcribe-2025-03-20": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -44444,6 +44761,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-07-23", "input_cost_per_audio_token": 1e-05, "input_cost_per_image": 8e-07, "input_cost_per_token": 6e-07, @@ -44524,6 +44842,7 @@ "supports_audio_input": true }, "sora-2": { + "deprecation_date": "2026-09-24", "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.1, @@ -44537,6 +44856,7 @@ ] }, "sora-2-pro": { + "deprecation_date": "2026-09-24", "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.3, @@ -44564,6 +44884,7 @@ }, "chatgpt-image-latest": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-12-01", "input_cost_per_image_token": 1e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -45714,6 +46035,7 @@ "supports_response_schema": true }, "snowflake/claude-sonnet-4-6": { + "supports_adaptive_thinking": true, "max_tokens": 16384, "max_input_tokens": 200000, "max_output_tokens": 16384, @@ -46135,40 +46457,6 @@ "supports_tool_choice": true, "supports_vision": false }, - "darkbloom/gemma-4-26b": { - "input_cost_per_token": 3e-08, - "litellm_provider": "darkbloom", - "max_input_tokens": 131072, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 1.65e-07, - "source": "https://www.darkbloom.dev/", - "supported_endpoints": [ - "/v1/chat/completions" - ], - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, - "darkbloom/gpt-oss-20b": { - "input_cost_per_token": 1.45e-08, - "litellm_provider": "darkbloom", - "max_input_tokens": 131072, - "max_output_tokens": 32768, - "max_tokens": 32768, - "mode": "chat", - "output_cost_per_token": 7e-08, - "source": "https://www.darkbloom.dev/", - "supported_endpoints": [ - "/v1/chat/completions" - ], - "supports_function_calling": true, - "supports_native_streaming": true, - "supports_system_messages": true, - "supports_tool_choice": true - }, "deepseek/deepseek-v4-pro": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 3.625e-09, @@ -46325,6 +46613,336 @@ "supports_reasoning": false, "source": "https://pinstripes.io/pricing" }, + "darkbloom/gemma-4-26b": { + "input_cost_per_token": 3e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.65e-07, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "darkbloom/gpt-oss-20b": { + "input_cost_per_token": 1.45e-08, + "litellm_provider": "darkbloom", + "max_input_tokens": 131072, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 7e-08, + "source": "https://www.darkbloom.dev/", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "xai/grok-4.20-0309-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true + }, + "xai/grok-4.20-multi-agent-0309": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true + }, + "xai/grok-build-0.1": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "xai", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true, + "supports_vision": true + }, + "gpt-transcribe": { + "input_cost_per_second": 7.5e-05, + "litellm_provider": "openai", + "mode": "audio_transcription", + "source": "https://platform.openai.com/docs/models/gpt-transcribe", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "gpt-live-transcribe": { + "input_cost_per_second": 0.0002833333333333333, + "litellm_provider": "openai", + "mode": "audio_transcription", + "source": "https://platform.openai.com/docs/models/gpt-live-transcribe", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "gpt-realtime-translate": { + "input_cost_per_second": 0.0005666666666666667, + "litellm_provider": "openai", + "max_input_tokens": 16000, + "max_output_tokens": 2000, + "max_tokens": 2000, + "mode": "realtime", + "source": "https://platform.openai.com/docs/models/gpt-realtime-translate", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, + "claude-mythos-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "prompt_cache_min_tokens": 512, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://docs.claude.com/en/docs/about-claude/models/overview", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "claude-mythos-preview": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "prompt_cache_min_tokens": 512, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://docs.claude.com/en/docs/about-claude/models/overview", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "gemini/gemini-robotics-er-2-streaming-preview": { + "input_cost_per_audio_token": 2e-06, + "input_cost_per_token": 2e-06, + "litellm_provider": "gemini", + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "mistral/mistral-small-2603": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/labs-leanstral-1-5": { + "input_cost_per_token": 0.0, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.mistral.ai/models/model-cards/leanstral-1-5", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/mistral-moderation-2603": { + "input_cost_per_token": 0.0, + "litellm_provider": "mistral", + "max_input_tokens": 131072, + "mode": "moderation", + "output_cost_per_token": 0.0, + "source": "https://docs.mistral.ai/models/model-cards/mistral-moderation-26-03" + }, + "mistral/voxtral-mini-2602": { + "input_cost_per_second": 5e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-transcribe-realtime-2602": { + "input_cost_per_second": 0.0001, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-realtime-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-tts-2603": { + "litellm_provider": "mistral", + "mode": "audio_speech", + "output_cost_per_character": 1.6e-05, + "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_output": true + }, "fallback_generalizations": { "rules": [ { diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index 0896c344f05..a9a3367cd93 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -15,6 +15,8 @@ from collections.abc import Mapping, Sequence from typing import Any, Final, NamedTuple, Optional, Protocol, Union, runtime_checkable if typing.TYPE_CHECKING: + from collections.abc import Awaitable, Callable + from fastapi import Request from mcp.client.session import ClientSession from mcp.shared.context import RequestContext @@ -28,8 +30,9 @@ if typing.TYPE_CHECKING: ToolUseContent, ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.utils import ProxyLogging + from litellm.types.utils import ModelResponse from fastapi import HTTPException from pydantic import TypeAdapter @@ -1016,7 +1019,7 @@ async def _run_budget_checks( general_settings=general_settings or {}, route="/chat/completions", llm_router=_llm_router, - proxy_logging_obj=typing.cast("ProxyLogging", _proxy_logging_obj), + proxy_logging_obj=_proxy_logging_obj, valid_token=user_api_key_auth, request=dummy_request, ) @@ -1176,15 +1179,19 @@ async def _build_completion_kwargs( ) +class _AcompletionCall(NamedTuple): + fn: "Callable[..., Awaitable[ModelResponse | CustomStreamWrapper]]" + + async def _run_guardrails_and_call_llm( - completion_kwargs: dict[str, Any], + completion_kwargs: dict[str, object], user_api_key_auth: "UserAPIKeyAuth", ) -> Any: try: from litellm.proxy.proxy_server import proxy_logging_obj as _plo if _plo is not None: - completion_kwargs = await typing.cast("ProxyLogging", _plo).pre_call_hook( + completion_kwargs = await _plo.pre_call_hook( user_api_key_dict=user_api_key_auth, data=completion_kwargs, call_type="acompletion", @@ -1204,10 +1211,10 @@ async def _run_guardrails_and_call_llm( from litellm.proxy.proxy_server import llm_router if llm_router is not None: - return await llm_router.acompletion(**completion_kwargs) - return await litellm.acompletion(**completion_kwargs) + return await _AcompletionCall(fn=llm_router.acompletion).fn(**completion_kwargs) + return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs) except ImportError: - return await litellm.acompletion(**completion_kwargs) + return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs) async def handle_sampling_create_message( diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index 36ffc819774..41359d44b27 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -3,8 +3,11 @@ Per-feature OpenAPI snapshot for lazy-loaded routers. The committed JSON is generated by `python -m litellm.proxy._lazy_openapi_snapshot` and consumed at runtime so /openapi.json can show full route info for unloaded -features without importing them. CI verifies the file is current and surfaces -any drift as a neutral check. +features without importing them. No CI job regenerates this file; drift surfaces +only indirectly through check-ui-api-types.yml, which rebuilds schema.d.ts from +app.openapi() with the committed snapshot injected. After changing any lazily +loaded route or this generator, rerun the module and commit the JSON, then run +`npm run gen:api` in ui/litellm-dashboard and commit schema.d.ts. """ import json @@ -89,8 +92,6 @@ def generate_snapshot() -> dict[str, dict]: from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids for feat in LAZY_FEATURES: - if feat.module_path in sys.modules: - continue try: module = importlib.import_module(feat.module_path) feat.register_fn(app, module) @@ -100,7 +101,7 @@ def generate_snapshot() -> dict[str, dict]: fragments: Final[dict[str, dict]] = {} used_operation_ids: Final[set[str]] = set() for feat in LAZY_FEATURES: - feat_routes = [r for r in app.routes if any(getattr(r, "path", "").startswith(p) for p in feat.path_prefixes)] + feat_routes = [r for r in app.routes if feat.matches(getattr(r, "path", ""))] if not feat_routes: continue _stabilize_multi_method_route_ids(feat_routes) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fa89df39c5f..dc4f17c7b31 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,7 +1,7 @@ import enum import json import os -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -11,6 +11,7 @@ from pydantic import ( ConfigDict, Field, Json, + PositiveInt, field_validator, model_validator, ) @@ -1102,6 +1103,8 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): class KeyRequestBase(GenerateRequestBase): key: str | None = None + default_estimated_output_tokens: PositiveInt | None = None + default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None budget_id: str | None = None tags: list[str] | None = None disable_global_guardrails: bool | None = None @@ -1819,6 +1822,8 @@ class NewTeamRequest(TeamBase): ) model_tpm_limit: dict[str, int] | None = None + default_estimated_output_tokens: PositiveInt | None = None + default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None mcp_rpm_limit: dict[str, int] | None = None team_member_budget: float | None = None # allow user to set a budget for all team members team_member_rpm_limit: int | None = None # allow user to set RPM limit for all team members @@ -1883,6 +1888,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): prompts: list[str] | None = None model_rpm_limit: dict[str, int] | None = None model_tpm_limit: dict[str, int] | None = None + default_estimated_output_tokens: PositiveInt | None = None + default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None mcp_rpm_limit: dict[str, int] | None = None allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None enforced_batch_output_expires_after: dict | None = None @@ -4018,6 +4025,13 @@ class LitellmDataForBackendLLMCall(TypedDict, total=False): stream_timeout: float | None user: str | None num_retries: int | None + # True when the effective timeout came from a caller-controlled source (the + # `x-litellm-timeout`/`x-litellm-stream-timeout` headers, or a `timeout`/`request_timeout`/ + # `stream_timeout` field in the request body) rather than deployment config, so a + # deliberately tiny value isn't treated as a deployment health signal (see + # cooldown_handlers._trigger_cooldown_for_failed_deployment). + client_side_timeout: bool + keepalive_seconds: float | None class LitellmMetadataFromRequestHeaders(TypedDict, total=False): @@ -4097,6 +4111,8 @@ class PassThroughEndpointLoggingTypedDict(TypedDict): LiteLLM_ManagementEndpoint_MetadataFields: Final = [ "model_rpm_limit", "model_tpm_limit", + "default_estimated_output_tokens", + "default_estimated_output_tokens_per_model", "mcp_rpm_limit", "tag_rpm_limit", "rpm_limit_type", diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 27780aeb994..497a39faf73 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -11,7 +11,7 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM """ import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping from copy import deepcopy from typing import TYPE_CHECKING, Any, Final from urllib.parse import urlparse @@ -36,7 +36,7 @@ from litellm.proxy.agent_endpoints.databricks_oauth import ( ) from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.utils import get_custom_url +from litellm.proxy.utils import ProxyLogging, get_custom_url from litellm.types.utils import all_litellm_params if TYPE_CHECKING: @@ -46,7 +46,7 @@ if TYPE_CHECKING: router: Final = APIRouter() -_PASCAL_TO_WIRE: Final[dict[str, str]] = { +_PASCAL_TO_WIRE: Final[Mapping[str, str]] = { "SendMessage": "message/send", "SendStreamingMessage": "message/stream", "GetTask": "tasks/get", @@ -118,9 +118,9 @@ def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> dict[str, str def _forwarding_headers( user_api_key_dict: UserAPIKeyAuth, - request_data: dict[str, Any], - agent_extra_headers: dict[str, str] | None, -) -> dict[str, str] | None: + request_data: Mapping[str, object], + agent_extra_headers: Mapping[str, str] | None, +) -> Mapping[str, str] | None: sanitized: Final = ( {k: v for k, v in agent_extra_headers.items() if not k.lower().startswith("x-litellm-")} if agent_extra_headers @@ -136,7 +136,7 @@ def _forwarding_headers( def _jsonrpc_error( - request_id: Any | None, + request_id: object, code: int, message: str, status_code: int = 400, @@ -162,7 +162,7 @@ def _get_agent(agent_id: str): return agent -def _enforce_inbound_trace_id(agent: Any, request: Request) -> None: +def _enforce_inbound_trace_id(agent: "AgentResponse", request: Request) -> None: """Raise 400 if agent requires x-litellm-trace-id on inbound calls and it is missing.""" agent_litellm_params: Final = agent.litellm_params or {} if not agent_litellm_params.get("require_trace_id_on_calls_to_agent"): @@ -181,8 +181,8 @@ def _enforce_inbound_trace_id(agent: Any, request: Request) -> None: async def _forward_jsonrpc( agent_url: str, - body: dict[str, Any], - extra_headers: dict[str, str] | None = None, + body: dict[str, object], + extra_headers: Mapping[str, str] | None = None, ) -> dict[str, Any]: from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider @@ -205,11 +205,11 @@ async def _forward_jsonrpc( async def _a2a_sse_event_source( agent_url: str, - body: dict[str, Any], - request_id: Any | None = None, - extra_headers: dict[str, str] | None = None, + body: Mapping[str, object], + request_id: str | int | None = None, + extra_headers: Mapping[str, str] | None = None, served_version: A2AVersion = "0.3", -) -> AsyncGenerator[dict, None]: +) -> AsyncGenerator[Mapping[str, object], None]: """Stream an upstream A2A SSE response as parsed JSON-RPC event dicts. Upstream HTTP/JSON-RPC errors are surfaced as a single JSON-RPC error event @@ -234,7 +234,7 @@ async def _a2a_sse_event_source( try: if not resp.is_success: error_body: Final = await resp.aread() - error_event: dict[str, Any] | None = None + error_event: Mapping[str, object] | None = None try: parsed: Final = json.loads(error_body) if isinstance(parsed, dict) and "error" in parsed: @@ -267,12 +267,12 @@ async def _a2a_sse_event_source( async def _forward_jsonrpc_sse( agent_url: str, - body: dict[str, Any], - request_id: Any | None = None, - extra_headers: dict[str, str] | None = None, - proxy_logging_obj: Any | None = None, - user_api_key_dict: Any | None = None, - request_data: dict[str, Any] | None = None, + body: Mapping[str, object], + request_id: str | int | None = None, + extra_headers: Mapping[str, str] | None = None, + proxy_logging_obj: ProxyLogging | None = None, + user_api_key_dict: UserAPIKeyAuth | None = None, + request_data: dict[str, object] | None = None, served_version: A2AVersion = "0.3", ) -> StreamingResponse: event_source: Final = _a2a_sse_event_source( @@ -283,10 +283,10 @@ async def _forward_jsonrpc_sse( served_version=served_version, ) - def _serialize_chunk(chunk: Any) -> str: + def _serialize_chunk(chunk: object) -> str: return f"data: {json.dumps(chunk)}\n\n" - def _serialize_error(proxy_exc: Any) -> str: + def _serialize_error(proxy_exc: object) -> str: return ( "data: " + json.dumps( @@ -331,17 +331,17 @@ async def _forward_jsonrpc_sse( async def _handle_stream_message( api_base: str | None, - request_id: Any, - params: dict[str, Any], - litellm_params: dict[str, Any] | None = None, + request_id: str | int, + params: dict[str, object], + litellm_params: dict[str, object] | None = None, agent_id: str | None = None, - metadata: dict[str, Any] | None = None, - proxy_server_request: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, + proxy_server_request: dict[str, object] | None = None, *, agent_extra_headers: dict[str, str] | None = None, user_api_key_dict: UserAPIKeyAuth | None = None, - request_data: dict[str, Any] | None = None, - proxy_logging_obj: Any | None = None, + request_data: dict[str, object] | None = None, + proxy_logging_obj: ProxyLogging | None = None, served_version: A2AVersion = "0.3", ) -> StreamingResponse: """Handle message/stream method via SDK functions. @@ -430,7 +430,7 @@ async def _handle_stream_message( obj = normalize_stream_event(obj, served_version, request_id=request_id) return json.dumps(obj) + "\n" - def _ndjson_error(proxy_exc: Any) -> str: + def _ndjson_error(proxy_exc: object) -> str: return ( json.dumps( { @@ -669,7 +669,7 @@ async def invoke_agent_a2a( agent_name: Final = agent_card_params.get("name", agent_id) # Get litellm_params (may include custom_llm_provider for completion bridge) - litellm_params = agent.litellm_params or {} + litellm_params: dict[str, object] = agent.litellm_params or {} custom_llm_provider: Final = litellm_params.get("custom_llm_provider") # Hand the authenticated key hash to the completion bridge so provider @@ -725,7 +725,7 @@ async def invoke_agent_a2a( request_data = data # Build merged headers for the backend agent - static_headers: Final[dict[str, str]] = dict(agent.static_headers or {}) + static_headers: Final[Mapping[str, str]] = dict(agent.static_headers or {}) raw_headers: Final = dict(request.headers) normalized: Final = {k.lower(): v for k, v in raw_headers.items()} @@ -893,7 +893,7 @@ async def invoke_agent_a2a( detail="Push notification URL must be a string", ) _validate_push_notification_url(callback_url) - forward_body = { + forward_body: dict[str, object] = { "jsonrpc": "2.0", "id": request_id, "method": method, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index b4e634b5eb1..3bfae4633c1 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1,12 +1,13 @@ import os import re import sys -from collections.abc import Iterator, Mapping +from collections.abc import Collection, Iterator, Mapping from functools import lru_cache from logging import Logger -from typing import Any, Final +from typing import Any, Final, Protocol from fastapi import HTTPException, Request, status +from pydantic import PositiveInt, TypeAdapter, ValidationError import litellm from litellm import Router, provider_list @@ -999,6 +1000,167 @@ def get_key_model_tpm_limit( return None +ESTIMATED_OUTPUT_TOKENS_FIELD: Final = "default_estimated_output_tokens" +ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD: Final = "default_estimated_output_tokens_per_model" +ESTIMATED_OUTPUT_TOKENS_METADATA_FIELDS: Final = frozenset( + {ESTIMATED_OUTPUT_TOKENS_FIELD, ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD} +) + +_ESTIMATED_OUTPUT_TOKENS_ADAPTER: Final = TypeAdapter(PositiveInt) +_ESTIMATED_OUTPUT_TOKENS_PER_MODEL_ADAPTER: Final = TypeAdapter(Mapping[str, PositiveInt]) + + +def _validated_output_token_estimate(raw: object) -> int | None: + """Coerce one declared estimate to a positive int, or ignore it.""" + if raw is None: + return None + try: + return _ESTIMATED_OUTPUT_TOKENS_ADAPTER.validate_python(raw) + except ValidationError as validation_error: + verbose_proxy_logger.warning( + "Ignoring malformed %s in metadata: %s", + ESTIMATED_OUTPUT_TOKENS_FIELD, + validation_error, + ) + return None + + +def _validated_output_token_estimates_per_model(raw: object) -> Mapping[str, int] | None: + """Coerce a declared per-model estimate map, or ignore it.""" + if raw is None: + return None + try: + return _ESTIMATED_OUTPUT_TOKENS_PER_MODEL_ADAPTER.validate_python(raw) + except ValidationError as validation_error: + verbose_proxy_logger.warning( + "Ignoring malformed %s in metadata: %s", + ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD, + validation_error, + ) + return None + + +def _estimated_output_tokens_from_metadata( + metadata: Mapping[str, Any] | None, + model_name: str | None, +) -> int | None: + """Resolve the per-model, then global, estimate out of one metadata blob. + + The two fields are validated independently so a malformed per-model map + cannot discard a valid global estimate, or the other way round. + """ + if not metadata or ESTIMATED_OUTPUT_TOKENS_METADATA_FIELDS.isdisjoint(metadata): + return None + + if model_name is not None: + per_model: Final = _validated_output_token_estimates_per_model( + metadata.get(ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD) + ) + per_model_estimate: Final = per_model.get(model_name) if per_model is not None else None + if per_model_estimate is not None: + return per_model_estimate + + return _validated_output_token_estimate(metadata.get(ESTIMATED_OUTPUT_TOKENS_FIELD)) + + +def get_estimated_output_tokens( + user_api_key_dict: UserAPIKeyAuth, + model_name: str | None = None, +) -> int | None: + """Resolve the operator-declared output-token estimate for TPM reservation. + + Priority order (returns first found): + 1. Key metadata ``default_estimated_output_tokens_per_model[model_name]`` + 2. Key metadata ``default_estimated_output_tokens`` + 3. Team metadata ``default_estimated_output_tokens_per_model[model_name]`` + 4. Team metadata ``default_estimated_output_tokens`` + + Returns ``None`` when nothing is configured, which leaves the static + heuristic floor in place. + """ + key_estimate: Final = _estimated_output_tokens_from_metadata(user_api_key_dict.metadata, model_name) + if key_estimate is not None: + return key_estimate + return _estimated_output_tokens_from_metadata(user_api_key_dict.team_metadata, model_name) + + +class OutputTokenEstimateRequest(Protocol): + """The shape of any management request that can carry an output-token estimate. + + Read-only members: the gate inspects a request, it never writes one back. + """ + + @property + def metadata(self) -> Mapping[str, object] | None: ... + + @property + def default_estimated_output_tokens(self) -> int | None: ... + + @property + def default_estimated_output_tokens_per_model(self) -> Mapping[str, int] | None: ... + + @property + def model_fields_set(self) -> Collection[str]: ... + + +def _requested_output_token_estimates( + data: OutputTokenEstimateRequest, + existing_metadata: Mapping[str, object], +) -> tuple[object, object]: + """The output-token estimates this request would leave stored on the entity. + + Mirrors how the management endpoints merge metadata: a supplied ``metadata`` + replaces the stored blob wholesale, an omitted one preserves it, and the + dedicated top-level fields overlay whatever survives. Both sources are read + because the same declaration reaches the same stored field either way. + """ + base: Final[Mapping[str, object]] = ( + (data.metadata or {}) if "metadata" in data.model_fields_set else existing_metadata + ) + return ( + data.default_estimated_output_tokens + if data.default_estimated_output_tokens is not None + else base.get(ESTIMATED_OUTPUT_TOKENS_FIELD), + data.default_estimated_output_tokens_per_model + if data.default_estimated_output_tokens_per_model is not None + else base.get(ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD), + ) + + +def enforce_output_token_estimates_are_admin_only( + data: OutputTokenEstimateRequest, + existing_metadata: Mapping[str, object] | None, + user_api_key_dict: UserAPIKeyAuth, + entity: Literal["key", "team"], +) -> None: + """Only a proxy admin may change what a key or team declares its models emit. + + That declaration is what the TPM limiter reserves for a request omitting + ``max_tokens``, so lowering or clearing it under-reserves against every + window the request is charged against, including the team and organization + ones the writer may not own. A key's metadata is writable by its holder and + a team's by its team admin, so neither is a trustworthy source for a value + that weakens a limit set above them. Gated on the resulting value rather + than on presence, so a form resending the stored declaration stays a no-op. + """ + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + stored: Final[Mapping[str, object]] = existing_metadata or {} + if _requested_output_token_estimates(data, stored) == ( + stored.get(ESTIMATED_OUTPUT_TOKENS_FIELD), + stored.get(ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD), + ): + return + raise HTTPException( + status_code=403, + detail={ + "error": f"Only proxy admins can set {ESTIMATED_OUTPUT_TOKENS_FIELD} or " + f"{ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD} on a {entity}. They decide how many output tokens " + "the rate limiter reserves for a request that omits max_tokens." + }, + ) + + def get_model_rate_limit_from_metadata( user_api_key_dict: UserAPIKeyAuth, metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], @@ -1141,18 +1303,27 @@ def is_pass_through_provider_route(route: str) -> bool: return False -def _has_user_setup_sso(): +def _has_user_setup_sso() -> bool: """ - Check if the user has set up single sign-on (SSO) by verifying the presence of Microsoft client ID, Google client ID or generic client ID and UI username environment variables. - Returns a boolean indicating whether SSO has been set up. + Check if the user has set up single sign-on (SSO). + + Covers OAuth providers (Microsoft, Google, generic) and SAML IdP metadata. + Used by UI discovery (``sso_configured``) so the login button enables when + any supported SSO path is configured — including SAML-only setups. """ microsoft_client_id: Final = os.getenv("MICROSOFT_CLIENT_ID", None) google_client_id: Final = os.getenv("GOOGLE_CLIENT_ID", None) generic_client_id: Final = os.getenv("GENERIC_CLIENT_ID", None) + saml_idp_metadata_url: Final = os.getenv("SAML_IDP_METADATA_URL", None) + saml_idp_metadata_xml: Final = os.getenv("SAML_IDP_METADATA_XML", None) - sso_setup = (microsoft_client_id is not None) or (google_client_id is not None) or (generic_client_id is not None) - - return sso_setup + return ( + microsoft_client_id is not None + or google_client_id is not None + or generic_client_id is not None + or bool(saml_idp_metadata_url) + or bool(saml_idp_metadata_xml) + ) def get_customer_user_header_from_mapping(user_id_mapping) -> list | None: diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index b7d936401fc..1cac515f9f2 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -11,6 +11,7 @@ import click import requests from rich.console import Console from rich.table import Table +from typing_extensions import NotRequired, TypedDict from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh @@ -18,6 +19,57 @@ from litellm.litellm_core_utils.cli_token_utils import is_cli_token_fresh from .private_json import write_private_json +class CliTokenData(TypedDict): + base_url: str + key: str + user_id: str + user_email: str + user_role: str + auth_header_name: str + jwt_token: str + timestamp: float + + +class CliTeam(TypedDict, total=False): + team_id: str | None + team_alias: str | None + models: list[str] + max_budget: float | None + + +class CliContextObj(TypedDict): + base_url: str + base_url_explicit: NotRequired[bool] + + +class CliPollData(TypedDict, total=False): + status: str + key: str + user_id: str + teams: list[str] + team_details: object + requires_team_selection: bool + team_id: str + + +class CliPollRequestKwargs(TypedDict, total=False): + timeout: int + headers: dict[str, str] + + +class CliSsoStartData(TypedDict): + login_id: str + poll_secret: str + user_code: str + + +class CliAuthResult(TypedDict): + api_key: str + user_id: str | None + teams: list[str] + team_id: str | None + + # Token storage utilities def get_token_file_path() -> str: """Get the path to store the authentication token""" @@ -27,12 +79,12 @@ def get_token_file_path() -> str: return str(config_dir / "token.json") -def save_token(token_data: dict[str, Any]) -> None: +def save_token(token_data: CliTokenData) -> None: """Save token data to file""" write_private_json(get_token_file_path(), token_data) -def load_token() -> dict[str, Any] | None: +def load_token() -> CliTokenData | None: """Load token data from file""" token_file: Final = get_token_file_path() if not os.path.exists(token_file): @@ -65,7 +117,7 @@ def get_stored_api_key(expected_base_url: str | None = None) -> str | None: # Team selection utilities -def display_teams_table(teams: list[dict[str, Any]]) -> None: +def display_teams_table(teams: list[CliTeam]) -> None: """Display teams in a formatted table""" console: Final = Console() @@ -165,7 +217,7 @@ def display_interactive_team_selection(teams: list[dict[str, Any]], selected_ind for i, team in enumerate(teams): team_alias = team.get("team_alias") or "N/A" team_id = team.get("team_id", "N/A") - models = team.get("models", []) + models: list[str] = team.get("models", []) max_budget = team.get("max_budget") # Format models list @@ -249,10 +301,11 @@ def prompt_team_selection_fallback( while True: try: - choice = click.prompt( + prompt_response: str = click.prompt( "\nSelect a team by entering the index number (or 'skip' to continue without a team)", type=str, - ).strip() + ) + choice = prompt_response.strip() if choice.lower() == "skip": return None @@ -275,7 +328,7 @@ def prompt_team_selection_fallback( def _response_error_detail(response: requests.Response) -> str | None: try: - body: Final = response.json() + body: Final[dict[str, object] | list[object] | str | int | float | bool | None] = response.json() except ValueError: return None detail: Final = body.get("detail") if isinstance(body, dict) else None @@ -309,15 +362,15 @@ def _poll_for_ready_data( other_status_log_every: int = 10, http_error_log_every: int = 10, connection_error_log_every: int = 10, -) -> dict[str, Any] | None: +) -> CliPollData | None: for attempt in range(total_timeout // poll_interval): try: - request_kwargs: dict[str, Any] = {"timeout": request_timeout} + request_kwargs: CliPollRequestKwargs = {"timeout": request_timeout} if headers is not None: request_kwargs["headers"] = headers response = requests.get(url, **request_kwargs) if response.status_code == 200: - data = response.json() + data: CliPollData = response.json() status = data.get("status") if status == "ready": return data @@ -341,7 +394,7 @@ def _poll_for_ready_data( return None -def _normalize_teams(teams, team_details): +def _normalize_teams(teams: object, team_details: object) -> list[CliTeam]: """If team_details are a Args: @@ -365,7 +418,7 @@ def _normalize_teams(teams, team_details): return [] -def _start_cli_sso_flow(base_url: str) -> dict[str, Any]: +def _start_cli_sso_flow(base_url: str) -> CliSsoStartData: start_url: Final = f"{base_url}/sso/cli/start" try: response: Final = requests.post(start_url, timeout=10) @@ -389,7 +442,7 @@ def _start_cli_sso_flow(base_url: str) -> dict[str, Any]: ) try: - data: Final = response.json() + data: Final[CliSsoStartData] = response.json() except ValueError: content_type: Final = response.headers.get("content-type", "unknown") raise ValueError( @@ -398,7 +451,7 @@ def _start_cli_sso_flow(base_url: str) -> dict[str, Any]: f"Response starts with: {response.text[:200]!r}" ) - required_fields: Final = ("login_id", "poll_secret", "user_code") + required_fields: Final[tuple[str, ...]] = ("login_id", "poll_secret", "user_code") missing_fields: Final = tuple(field for field in required_fields if not isinstance(data.get(field), str)) if missing_fields: raise ValueError( @@ -412,7 +465,7 @@ def _get_cli_sso_poll_headers(poll_secret: str) -> dict[str, str]: return {"x-litellm-cli-poll-secret": poll_secret} -def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> dict | None: +def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> CliAuthResult | None: """ Poll the server for authentication completion and handle team selection. @@ -431,7 +484,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> di teams = data.get("teams", []) team_details: Final = data.get("team_details") user_id = data.get("user_id") - normalized_teams: Final[list[dict[str, Any]]] = _normalize_teams(teams, team_details) + normalized_teams: Final[list[CliTeam]] = _normalize_teams(teams, team_details) if not normalized_teams: click.echo("Warning: No teams available for selection.") return None @@ -478,7 +531,7 @@ def _poll_for_authentication(base_url: str, key_id: str, poll_secret: str) -> di def _handle_team_selection_during_polling( - base_url: str, key_id: str, poll_secret: str, teams: list[dict[str, Any]] + base_url: str, key_id: str, poll_secret: str, teams: list[CliTeam] ) -> str | None: """ Handle team selection and re-poll with selected team_id. @@ -522,7 +575,7 @@ def _handle_team_selection_during_polling( return None -def _render_and_prompt_for_team_selection(teams: list[dict[str, Any]]) -> str | None: +def _render_and_prompt_for_team_selection(teams: list[CliTeam]) -> str | None: """Render teams table and prompt user for a team selection. Returns the selected team_id as a string, or None if selection was @@ -546,10 +599,11 @@ def _render_and_prompt_for_team_selection(teams: list[dict[str, Any]]) -> str | # Simple selection while True: try: - choice = click.prompt( + prompt_response: str = click.prompt( "\nSelect a team by entering the index number (or 'skip' to use first team)", type=str, - ).strip() + ) + choice = prompt_response.strip() if choice.lower() == "skip": # Default to the first team's ID if the user skips an @@ -582,7 +636,8 @@ def login(ctx: click.Context): from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER from litellm.proxy.client.cli.interface import show_commands - base_url: Final = ctx.obj["base_url"] + ctx_obj: Final[CliContextObj] = ctx.obj + base_url: Final = ctx_obj["base_url"] try: cli_sso_flow: Final = _start_cli_sso_flow(base_url=base_url) @@ -675,8 +730,9 @@ def print_token(ctx: click.Context): # explicitly pointed us at a server, trust whichever one `lite login` # actually issued this token for -- that's the whole point of not # needing a wrapper command. - if ctx.obj.get("base_url_explicit"): - base_url: Final = ctx.obj["base_url"] + ctx_obj: Final[CliContextObj] = ctx.obj + if ctx_obj.get("base_url_explicit"): + base_url: Final = ctx_obj["base_url"] if token_data.get("base_url") != base_url.rstrip("/"): click.echo("Not authenticated for this server. Run 'lite login'.", err=True) sys.exit(1) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 8830970f96f..bf760a92d88 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1,14 +1,21 @@ import asyncio import json import time -from collections.abc import Callable, Sequence +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timezone -from typing import Final, Literal, Protocol, TypeVar +from types import MappingProxyType +from typing import Final, Literal, Protocol, TypeVar, assert_never import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache -from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME +from litellm.constants import ( + GLOBAL_PROXY_SPEND_CACHE_KEY, + LITELLM_PROXY_BUDGET_NAME, + RESET_BUDGET_JOB_BATCH_SIZE, + RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN, +) from litellm.proxy._types import ( LiteLLM_BudgetTableFull, LiteLLM_EndUserTable, @@ -30,7 +37,10 @@ from litellm.repositories.table_repositories import ( TeamMembershipRepository, ) from litellm.repositories.team_repository import TeamRepository -from litellm.repositories.unit_of_work import spend_reset_unit_of_work +from litellm.repositories.unit_of_work import ( + budget_cascade_unit_of_work, + spend_reset_unit_of_work, +) from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) @@ -38,6 +48,9 @@ from litellm.types.services import ServiceTypes _RowT = TypeVar("_RowT") +_LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}}) +_SPENT_ROWS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"spend": {"gt": 0}}) + class _TeamMembershipRow(Protocol): @property @@ -62,39 +75,130 @@ class _TagRow(Protocol): def tag_name(self) -> str: ... +class _EndUserRow(Protocol): + @property + def user_id(self) -> str: ... + + def _team_membership_counter_key(row: _TeamMembershipRow) -> str: return f"spend:team_member:{row.user_id}:{row.team_id}" -def _team_membership_cache_key(row: _TeamMembershipRow) -> str: - return f"{row.team_id}_{row.user_id}" +def _team_membership_cache_keys(row: _TeamMembershipRow) -> tuple[str, ...]: + return (f"{row.team_id}_{row.user_id}",) def _key_counter_key(row: _KeyRow) -> str: return f"spend:key:{row.token}" -def _key_cache_key(row: _KeyRow) -> str: - return row.token +def _key_cache_keys(row: _KeyRow) -> tuple[str, ...]: + return (row.token,) def _org_counter_key(row: _OrgRow) -> str: return f"spend:org:{row.organization_id}" -def _org_cache_keys(row: _OrgRow) -> Sequence[str]: - return [ +def _org_cache_keys(row: _OrgRow) -> tuple[str, ...]: + return ( f"org_id:{row.organization_id}", f"org_id:{row.organization_id}:with_budget", - ] + ) def _tag_counter_key(row: _TagRow) -> str: return f"spend:tag:{row.tag_name}" -def _tag_cache_key(row: _TagRow) -> str: - return f"tag:{row.tag_name}" +def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]: + return (f"tag:{row.tag_name}",) + + +def _budget_link_where( + budget_ids: Sequence[str], + extra: Mapping[str, object] = MappingProxyType({}), +) -> dict[str, object]: + return {"budget_id": {"in": list(budget_ids)}, **extra} + + +@dataclass(frozen=True, slots=True) +class _BudgetCascade: + """Everything one budget-tier reset touches, resolved before any write.""" + + budgets: tuple[LiteLLM_BudgetTableFull, ...] = () + budget_ids: tuple[str, ...] = () + budget_resets: tuple[tuple[str, datetime], ...] = () + endusers: tuple[_EndUserRow, ...] = () + counter_keys: tuple[str, ...] = () + cache_keys: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class _BudgetCascadeCommitted: + cascade: _BudgetCascade + advanced: int + + +@dataclass(frozen=True, slots=True) +class _BudgetCascadeFailed: + cascade: _BudgetCascade + error: Exception + + +_EMPTY_CASCADE: Final = _BudgetCascade() + + +@dataclass(frozen=True, slots=True) +class _ChunkOutcome: + """One chunk of a reset phase: rows read, and rows whose new budget_reset_at + cleared the due cutoff. Anything else is still due and would come straight + back on the next fetch, so it is not progress.""" + + fetched: int + advanced: int + + +_NO_PROGRESS: Final = _ChunkOutcome(fetched=0, advanced=0) + + +def _as_utc(moment: datetime) -> datetime: + return moment if moment.tzinfo is not None else moment.replace(tzinfo=timezone.utc) + + +def _count_advanced(reset_ats: Iterable[object], cutoff: datetime) -> int: + """How many rows the write actually moved past the due cutoff. + + A budget_duration of "0s" (or one the parser cannot read) resolves to the + current time, so the row is written and stays due. Counting it as progress + would re-read the same chunk until the per-run cap on every tick. + """ + utc_cutoff: Final = _as_utc(cutoff) + return sum(1 for reset_at in reset_ats if isinstance(reset_at, datetime) and _as_utc(reset_at) > utc_cutoff) + + +def _phase_is_drained(outcome: _ChunkOutcome) -> bool: + """A short chunk means the due rows ran out. A full chunk that advanced + nothing would be re-read unchanged forever, so it ends the phase too and + those rows wait for the next tick.""" + return outcome.fetched < RESET_BUDGET_JOB_BATCH_SIZE or outcome.advanced == 0 + + +async def _run_phase_in_chunks(process_chunk: Callable[[], Awaitable[_ChunkOutcome]]) -> None: + """Drive one reset phase a chunk at a time, capped so a single run cannot + spin unbounded: leftovers are picked up by the next tick.""" + for _ in range(RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN): + if _phase_is_drained(await process_chunk()): + return + + +def _budget_cascade_event_metadata(cascade: _BudgetCascade) -> dict[str, object]: + return { + "num_budgets_found": len(cascade.budgets), + "budgets_found": json.dumps(cascade.budgets, indent=4, default=str), + "num_endusers_found": len(cascade.endusers), + "endusers_found": json.dumps(cascade.endusers, indent=4, default=str), + } class ResetBudgetJob: @@ -122,21 +226,14 @@ class ResetBudgetJob: Updates db """ - if self.prisma_client is not None: - ### RESET KEY BUDGET ### - await self.reset_budget_for_litellm_keys() + if self.prisma_client is None: + return - ### RESET USER BUDGET ### - await self.reset_budget_for_litellm_users() - - ## Reset Team Budget - await self.reset_budget_for_litellm_teams() - - ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ### - await self.reset_budget_for_litellm_budget_table() - - ### RESET MULTI-WINDOW BUDGETS ### - await self.reset_budget_windows() + await self.reset_budget_for_litellm_keys() + await self.reset_budget_for_litellm_users() + await self.reset_budget_for_litellm_teams() + await self.reset_budget_for_litellm_budget_table() + await self.reset_budget_windows() @staticmethod async def _invalidate_spend_counter(counter_key: str) -> None: @@ -194,238 +291,195 @@ class ResetBudgetJob: e, ) - async def _cascade_reset_spend_for_budget_link( + async def _fetch_linked_rows( self, - budgets_to_reset: list[LiteLLM_BudgetTableFull], table: SpendLinkedTable[_RowT], - counter_key_fn: Callable[[_RowT], str], + where: Mapping[str, object], log_subject: str, - extra_where: dict[str, object] | None = None, - cache_key_fn: Callable[[_RowT], str | Sequence[str]] | None = None, - ): - """ - Generic cascade: zero spend on rows whose budget_id is in the reset set. + ) -> tuple[_RowT, ...]: + """Read the rows the cascade will zero, so their counters can be + invalidated once the transaction commits.""" + try: + return tuple(await table.find_many(where=where)) + except Exception as e: + verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e) + return () - ``cache_key_fn`` is optional: when provided, after the DB update each - matching row's entry or entries in ``user_api_key_cache`` are dropped so - cached spend cannot stay pinned above the zeroed DB row after a reset. + async def _collect_endusers_to_reset(self, budget_ids: Sequence[str]) -> tuple[_EndUserRow, ...]: + linked: Final[Sequence[_EndUserRow] | None] = await self.prisma_client.get_data( + table_name="enduser", + query_type="find_all", + budget_id_list=list(budget_ids), + ) + if litellm.max_end_user_budget_id is None or litellm.max_end_user_budget_id not in budget_ids: + return tuple(linked or ()) + return (*(linked or ()), *await self._get_endusers_with_no_budget_id()) + + async def _collect_budget_cascade(self, budgets_to_reset: Sequence[LiteLLM_BudgetTableFull]) -> _BudgetCascade: + """Resolve every row the expiring budget tiers gate, before any write. + + Keys carrying their own budget_duration are left out: they run on their + own schedule via reset_budget_for_litellm_keys(), so sweeping them here + would reset them twice. """ - budget_ids: Final = [b.budget_id for b in budgets_to_reset if b.budget_id is not None] + budget_ids: Final = tuple(b.budget_id for b in budgets_to_reset if b.budget_id is not None) if not budget_ids: + return _EMPTY_CASCADE + + team_memberships: Final[tuple[_TeamMembershipRow, ...]] = await self._fetch_linked_rows( + table=TeamMembershipRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids), + log_subject="team memberships", + ) + keys: Final[tuple[_KeyRow, ...]] = await self._fetch_linked_rows( + table=VerificationTokenRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _LINKED_KEYS_WHERE), + log_subject="keys", + ) + orgs: Final[tuple[_OrgRow, ...]] = await self._fetch_linked_rows( + table=OrganizationRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), + log_subject="orgs", + ) + tags: Final[tuple[_TagRow, ...]] = await self._fetch_linked_rows( + table=TagRepository(self.prisma_client).table, + where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE), + log_subject="tags", + ) + return _BudgetCascade( + budgets=tuple(budgets_to_reset), + budget_ids=budget_ids, + budget_resets=tuple( + ( + b.budget_id, + compute_budget_reset_at(budget_duration=b.budget_duration, settings=self.reset_settings), + ) + for b in budgets_to_reset + if b.budget_id is not None and b.budget_duration is not None + ), + endusers=await self._collect_endusers_to_reset(budget_ids), + counter_keys=( + *(_team_membership_counter_key(row) for row in team_memberships), + *(_key_counter_key(row) for row in keys), + *(_org_counter_key(row) for row in orgs), + *(_tag_counter_key(row) for row in tags), + ), + cache_keys=( + *(key for row in team_memberships for key in _team_membership_cache_keys(row)), + *(key for row in keys for key in _key_cache_keys(row)), + *(key for row in orgs for key in _org_cache_keys(row)), + *(key for row in tags for key in _tag_cache_keys(row)), + ), + ) + + async def _commit_budget_cascade(self, cascade: _BudgetCascade) -> None: + """Zero the gated spend and advance ``budget_reset_at`` in one transaction. + + Advancing the window on its own hides the tier from every later tick + while its dependents stay pinned at the cap for the whole window; + batching both means a mid-cascade failure persists nothing and the rows + stay due for the next run. + """ + if not cascade.budget_ids: return - where: Final[dict[str, object]] = {"budget_id": {"in": budget_ids}} - if extra_where: - where.update(extra_where) + enduser_ids: Final = tuple(row.user_id for row in cascade.endusers) + async with budget_cascade_unit_of_work(self.prisma_client.db.batch_) as uow: + uow.team_memberships.queue_spend_zero(where=_budget_link_where(cascade.budget_ids)) + uow.keys.queue_spend_zero(where=_budget_link_where(cascade.budget_ids, _LINKED_KEYS_WHERE)) + uow.organizations.queue_spend_zero(where=_budget_link_where(cascade.budget_ids, _SPENT_ROWS_WHERE)) + uow.tags.queue_spend_zero(where=_budget_link_where(cascade.budget_ids, _SPENT_ROWS_WHERE)) + if enduser_ids: + uow.endusers.queue_spend_zero(where={"user_id": {"in": list(enduser_ids)}}) + for budget_id, budget_reset_at in cascade.budget_resets: + uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) - try: - rows: Sequence[_RowT] = await table.find_many(where=where) - except Exception as e: - rows = () - verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e) - - update_result: Final = await table.update_many(where=where, data={"spend": 0}) - - for row in rows: - await self._invalidate_spend_counter(counter_key_fn(row)) - if cache_key_fn is not None: - cache_keys = cache_key_fn(row) - if isinstance(cache_keys, str): - cache_keys = [cache_keys] - for cache_key in cache_keys: - await self._invalidate_user_api_key_cache_entry(cache_key) - - return update_result - - async def reset_budget_for_litellm_team_members(self, budgets_to_reset: list[LiteLLM_BudgetTableFull]): - """ - Resets the budget for all LiteLLM Team Members if their budget has expired - """ - return await self._cascade_reset_spend_for_budget_link( - budgets_to_reset=budgets_to_reset, - table=TeamMembershipRepository(self.prisma_client).table, - counter_key_fn=_team_membership_counter_key, - log_subject="team memberships", - cache_key_fn=_team_membership_cache_key, - ) - - async def reset_budget_for_keys_linked_to_budgets(self, budgets_to_reset: list[LiteLLM_BudgetTableFull]): - """ - Resets the spend for keys linked to budget tiers that are being reset. - - Excludes keys with their own budget_duration; those are reset by - reset_budget_for_litellm_keys() to avoid double-resetting. - """ - return await self._cascade_reset_spend_for_budget_link( - budgets_to_reset=budgets_to_reset, - table=VerificationTokenRepository(self.prisma_client).table, - counter_key_fn=_key_counter_key, - log_subject="keys", - extra_where={"budget_duration": None, "spend": {"gt": 0}}, - cache_key_fn=_key_cache_key, - ) - - async def reset_budget_for_orgs_linked_to_budgets(self, budgets_to_reset: list[LiteLLM_BudgetTableFull]): - """ - Resets the spend for orgs linked to budget tiers that are being reset. - """ - return await self._cascade_reset_spend_for_budget_link( - budgets_to_reset=budgets_to_reset, - table=OrganizationRepository(self.prisma_client).table, - counter_key_fn=_org_counter_key, - log_subject="orgs", - extra_where={"spend": {"gt": 0}}, - cache_key_fn=_org_cache_keys, - ) - - async def reset_budget_for_tags_linked_to_budgets(self, budgets_to_reset: list[LiteLLM_BudgetTableFull]): - """ - Resets the spend for tags linked to budget tiers that are being reset. - - Also drops each tag's ``user_api_key_cache`` entry so the next - ``_tag_max_budget_check`` reloads the zeroed row from the DB. - ``SpendCounterReseed.from_db`` intentionally returns ``None`` for - tags, so the budget check falls back to the cached - ``LiteLLM_TagTable.spend`` once the spend counter expires; without - this invalidation, that stale ``.spend`` keeps the tag over-budget - indefinitely. - """ - return await self._cascade_reset_spend_for_budget_link( - budgets_to_reset=budgets_to_reset, - table=TagRepository(self.prisma_client).table, - counter_key_fn=_tag_counter_key, - log_subject="tags", - extra_where={"spend": {"gt": 0}}, - cache_key_fn=_tag_cache_key, - ) - - async def reset_budget_for_litellm_budget_table(self): - """ - Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired - The corresponding Budget duration is also updated. - """ + async def _invalidate_budget_cascade_caches(self, cascade: _BudgetCascade) -> None: + for counter_key in cascade.counter_keys: + await self._invalidate_spend_counter(counter_key) + for cache_key in cascade.cache_keys: + await self._invalidate_user_api_key_cache_entry(cache_key) + async def _reset_expired_budget_cascade(self) -> _BudgetCascadeCommitted | _BudgetCascadeFailed: now: Final = datetime.now(timezone.utc) - start_time: Final = time.time() - endusers_to_reset: list[LiteLLM_EndUserTable] | None = None - budgets_to_reset: list[LiteLLM_BudgetTableFull] | None = None - updated_endusers: Final[list[LiteLLM_EndUserTable]] = [] - failed_endusers: Final = [] try: - budgets_to_reset = await self.prisma_client.get_data( - table_name="budget", query_type="find_all", reset_at=now - ) - - if budgets_to_reset is not None and len(budgets_to_reset) > 0: - for budget in budgets_to_reset: - budget = await ResetBudgetJob._reset_budget_reset_at_date(budget, now, self.reset_settings) - - await self.prisma_client.update_data( - query_type="update_many", - data_list=budgets_to_reset, - table_name="budget", - ) - - budget_ids_to_reset = [budget.budget_id for budget in budgets_to_reset if budget.budget_id is not None] - - endusers_to_reset = await self.prisma_client.get_data( - table_name="enduser", - query_type="find_all", - budget_id_list=budget_ids_to_reset, - ) - - # Also reset end users with no budget_id (NULL) who use the - # default budget via litellm.max_end_user_budget_id. These - # users are enforced in-memory but never had budget_id - # persisted, so the query above misses them. - if litellm.max_end_user_budget_id is not None and litellm.max_end_user_budget_id in budget_ids_to_reset: - default_budget_endusers: Final = await self._get_endusers_with_no_budget_id() - if default_budget_endusers: - if endusers_to_reset is None: - endusers_to_reset = default_budget_endusers - else: - endusers_to_reset.extend(default_budget_endusers) - - await self.reset_budget_for_litellm_team_members(budgets_to_reset=budgets_to_reset) - - await self.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset) - - await self.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=budgets_to_reset) - - await self.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=budgets_to_reset) - - if endusers_to_reset is not None and len(endusers_to_reset) > 0: - for enduser in endusers_to_reset: - try: - updated_enduser = await ResetBudgetJob._reset_budget_for_enduser(enduser=enduser) - if updated_enduser is not None: - updated_endusers.append(updated_enduser) - else: - failed_endusers.append( - { - "enduser": enduser, - "error": "Returned None without exception", - } - ) - except Exception as e: - failed_endusers.append({"enduser": enduser, "error": str(e)}) - verbose_proxy_logger.exception("Failed to reset budget for enduser: %s", enduser) - - verbose_proxy_logger.debug( - "Updated users %s", - json.dumps(updated_endusers, indent=4, default=str), - ) - - await self.prisma_client.update_data( - query_type="update_many", - data_list=updated_endusers, - table_name="enduser", - ) - - end_time = time.time() - if len(failed_endusers) > 0: # If any endusers failed to reset - raise Exception( - f"Failed to reset {len(failed_endusers)} endusers: {json.dumps(failed_endusers, default=str)}" - ) - - asyncio.create_task( - self.proxy_logging_obj.service_logging_obj.async_service_success_hook( - service=ServiceTypes.RESET_BUDGET_JOB, - duration=end_time - start_time, - call_type="reset_budget_budget_table", - start_time=start_time, - end_time=end_time, - event_metadata={ - "num_budgets_found": (len(budgets_to_reset) if budgets_to_reset else 0), - "budgets_found": json.dumps(budgets_to_reset, indent=4, default=str), - "num_endusers_found": (len(endusers_to_reset) if endusers_to_reset else 0), - "endusers_found": json.dumps(endusers_to_reset, indent=4, default=str), - "num_endusers_updated": len(updated_endusers), - "endusers_updated": json.dumps(updated_endusers, indent=4, default=str), - "num_endusers_failed": len(failed_endusers), - "endusers_failed": json.dumps(failed_endusers, indent=4, default=str), - }, - ) + budgets_to_reset: Final[Sequence[LiteLLM_BudgetTableFull] | None] = await self.prisma_client.get_data( + table_name="budget", + query_type="find_all", + reset_at=now, + limit=RESET_BUDGET_JOB_BATCH_SIZE, ) + cascade: Final = await self._collect_budget_cascade(budgets_to_reset or ()) except Exception as e: - end_time = time.time() - asyncio.create_task( - self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( - service=ServiceTypes.RESET_BUDGET_JOB, - duration=end_time - start_time, - error=e, - call_type="reset_budget_endusers", - start_time=start_time, - end_time=end_time, - event_metadata={ - "num_budgets_found": (len(budgets_to_reset) if budgets_to_reset else 0), - "budgets_found": json.dumps(budgets_to_reset, indent=4, default=str), - "num_endusers_found": (len(endusers_to_reset) if endusers_to_reset else 0), - "endusers_found": json.dumps(endusers_to_reset, indent=4, default=str), - }, + return _BudgetCascadeFailed(cascade=_EMPTY_CASCADE, error=e) + + try: + await self._commit_budget_cascade(cascade) + except Exception as e: + return _BudgetCascadeFailed(cascade=cascade, error=e) + + await self._invalidate_budget_cascade_caches(cascade) + return _BudgetCascadeCommitted( + cascade=cascade, + advanced=_count_advanced( + (reset_at for _, reset_at in cascade.budget_resets), + cutoff=datetime.now(timezone.utc), + ), + ) + + async def reset_budget_for_litellm_budget_table(self) -> None: + """ + Resets the spend a budget tier gates (end users, team members, keys, + orgs, tags) and advances the tier's budget_reset_at, atomically. + + Caches are invalidated only after the transaction commits, so a failed + run cannot leave a zeroed counter in front of an un-reset DB row. + """ + await _run_phase_in_chunks(self._reset_budget_for_litellm_budget_table_chunk) + + async def _reset_budget_for_litellm_budget_table_chunk(self) -> _ChunkOutcome: + start_time: Final = time.time() + outcome: Final = await self._reset_expired_budget_cascade() + end_time: Final = time.time() + + match outcome: + case _BudgetCascadeCommitted(cascade=cascade, advanced=advanced): + asyncio.create_task( + self.proxy_logging_obj.service_logging_obj.async_service_success_hook( + service=ServiceTypes.RESET_BUDGET_JOB, + duration=end_time - start_time, + call_type="reset_budget_budget_table", + start_time=start_time, + end_time=end_time, + event_metadata={ + **_budget_cascade_event_metadata(cascade), + "num_endusers_updated": len(cascade.endusers), + "num_endusers_failed": 0, + }, + ) ) - ) - verbose_proxy_logger.exception("Failed to reset budget for endusers: %s", e) + return _ChunkOutcome(fetched=len(cascade.budgets), advanced=advanced) + case _BudgetCascadeFailed(cascade=cascade, error=error): + verbose_proxy_logger.exception( + "Failed to reset the budget table cascade (team member, enduser, org and tag spend, plus " + "budget_reset_at); nothing was committed and the budgets stay due for the next run: %s", + error, + exc_info=error, + ) + asyncio.create_task( + self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( + service=ServiceTypes.RESET_BUDGET_JOB, + duration=end_time - start_time, + error=error, + call_type="reset_budget_endusers", + start_time=start_time, + end_time=end_time, + event_metadata=_budget_cascade_event_metadata(cascade), + ) + ) + return _NO_PROGRESS + case _: + assert_never(outcome) async def _get_endusers_with_no_budget_id( self, @@ -486,18 +540,50 @@ class ResetBudgetJob: for t in updated_teams: uow.teams.queue_spend_reset(team_id=t.team_id, budget_reset_at=t.budget_reset_at) - async def reset_budget_for_litellm_keys(self): + def _emit_phase_failure( + self, + call_type: str, + error: Exception, + start_time: float, + end_time: float, + event_metadata: dict[str, object], + ) -> None: + """Report rows that could not be reset without failing the chunk: the + rows that did reset are already committed, and raising here would cost + the phase every remaining chunk this tick. + """ + verbose_proxy_logger.error("%s: %s", call_type, error) + asyncio.create_task( + self.proxy_logging_obj.service_logging_obj.async_service_failure_hook( + service=ServiceTypes.RESET_BUDGET_JOB, + duration=end_time - start_time, + error=error, + call_type=call_type, + start_time=start_time, + end_time=end_time, + event_metadata=event_metadata, + ) + ) + + async def reset_budget_for_litellm_keys(self) -> None: """ Resets the budget for all the litellm keys Catches Exceptions and logs them """ + await _run_phase_in_chunks(self._reset_budget_for_litellm_keys_chunk) + + async def _reset_budget_for_litellm_keys_chunk(self) -> _ChunkOutcome: now: Final = datetime.utcnow() start_time: Final = time.time() keys_to_reset: list[LiteLLM_VerificationToken] | None = None try: keys_to_reset = await self.prisma_client.get_data( - table_name="key", query_type="find_all", expires=now, reset_at=now + table_name="key", + query_type="find_all", + expires=now, + reset_at=now, + limit=RESET_BUDGET_JOB_BATCH_SIZE, ) verbose_proxy_logger.debug("Keys to reset %s", json.dumps(keys_to_reset, indent=4, default=str)) updated_keys: Final[list[LiteLLM_VerificationToken]] = [] @@ -528,8 +614,25 @@ class ResetBudgetJob: await self._invalidate_spend_counter(f"spend:key:{token}") end_time = time.time() - if len(failed_keys) > 0: # If any keys failed to reset - raise Exception(f"Failed to reset {len(failed_keys)} keys: {json.dumps(failed_keys, default=str)}") + outcome: Final = _ChunkOutcome( + fetched=len(keys_to_reset) if keys_to_reset else 0, + advanced=_count_advanced( + (k.budget_reset_at for k in updated_keys), + cutoff=datetime.now(timezone.utc), + ), + ) + if len(failed_keys) > 0: + self._emit_phase_failure( + call_type="reset_budget_keys", + error=Exception(f"Failed to reset {len(failed_keys)} keys: {json.dumps(failed_keys, default=str)}"), + start_time=start_time, + end_time=end_time, + event_metadata={ + "num_keys_found": len(keys_to_reset) if keys_to_reset else 0, + "keys_found": json.dumps(keys_to_reset, indent=4, default=str), + }, + ) + return outcome asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( @@ -565,16 +668,27 @@ class ResetBudgetJob: ) ) verbose_proxy_logger.exception("Failed to reset budget for keys: %s", e) + return _NO_PROGRESS + else: + return outcome - async def reset_budget_for_litellm_users(self): + async def reset_budget_for_litellm_users(self) -> None: """ Resets the budget for all LiteLLM Internal Users if their budget has expired """ + await _run_phase_in_chunks(self._reset_budget_for_litellm_users_chunk) + + async def _reset_budget_for_litellm_users_chunk(self) -> _ChunkOutcome: now: Final = datetime.utcnow() start_time: Final = time.time() users_to_reset: list[LiteLLM_UserTable] | None = None try: - users_to_reset = await self.prisma_client.get_data(table_name="user", query_type="find_all", reset_at=now) + users_to_reset = await self.prisma_client.get_data( + table_name="user", + query_type="find_all", + reset_at=now, + limit=RESET_BUDGET_JOB_BATCH_SIZE, + ) updated_users: Final[list[LiteLLM_UserTable]] = [] failed_users: Final = [] if users_to_reset is not None and len(users_to_reset) > 0: @@ -609,8 +723,27 @@ class ResetBudgetJob: await self._invalidate_global_proxy_spend_cache() end_time = time.time() - if len(failed_users) > 0: # If any users failed to reset - raise Exception(f"Failed to reset {len(failed_users)} users: {json.dumps(failed_users, default=str)}") + outcome: Final = _ChunkOutcome( + fetched=len(users_to_reset) if users_to_reset else 0, + advanced=_count_advanced( + (u.budget_reset_at for u in updated_users), + cutoff=datetime.now(timezone.utc), + ), + ) + if len(failed_users) > 0: + self._emit_phase_failure( + call_type="reset_budget_users", + error=Exception( + f"Failed to reset {len(failed_users)} users: {json.dumps(failed_users, default=str)}" + ), + start_time=start_time, + end_time=end_time, + event_metadata={ + "num_users_found": len(users_to_reset) if users_to_reset else 0, + "users_found": json.dumps(users_to_reset, indent=4, default=str), + }, + ) + return outcome asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( @@ -646,16 +779,27 @@ class ResetBudgetJob: ) ) verbose_proxy_logger.exception("Failed to reset budget for users: %s", e) + return _NO_PROGRESS + else: + return outcome - async def reset_budget_for_litellm_teams(self): + async def reset_budget_for_litellm_teams(self) -> None: """ Resets the budget for all LiteLLM Internal Teams if their budget has expired """ + await _run_phase_in_chunks(self._reset_budget_for_litellm_teams_chunk) + + async def _reset_budget_for_litellm_teams_chunk(self) -> _ChunkOutcome: now: Final = datetime.utcnow() start_time: Final = time.time() teams_to_reset: list[LiteLLM_TeamTable] | None = None try: - teams_to_reset = await self.prisma_client.get_data(table_name="team", query_type="find_all", reset_at=now) + teams_to_reset = await self.prisma_client.get_data( + table_name="team", + query_type="find_all", + reset_at=now, + limit=RESET_BUDGET_JOB_BATCH_SIZE, + ) updated_teams: Final[list[LiteLLM_TeamTable]] = [] failed_teams: Final = [] if teams_to_reset is not None and len(teams_to_reset) > 0: @@ -688,8 +832,27 @@ class ResetBudgetJob: await self._invalidate_spend_counter(f"spend:team:{team_id}") end_time = time.time() - if len(failed_teams) > 0: # If any teams failed to reset - raise Exception(f"Failed to reset {len(failed_teams)} teams: {json.dumps(failed_teams, default=str)}") + outcome: Final = _ChunkOutcome( + fetched=len(teams_to_reset) if teams_to_reset else 0, + advanced=_count_advanced( + (t.budget_reset_at for t in updated_teams), + cutoff=datetime.now(timezone.utc), + ), + ) + if len(failed_teams) > 0: + self._emit_phase_failure( + call_type="reset_budget_teams", + error=Exception( + f"Failed to reset {len(failed_teams)} teams: {json.dumps(failed_teams, default=str)}" + ), + start_time=start_time, + end_time=end_time, + event_metadata={ + "num_teams_found": len(teams_to_reset) if teams_to_reset else 0, + "teams_found": json.dumps(teams_to_reset, indent=4, default=str), + }, + ) + return outcome asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( @@ -725,6 +888,9 @@ class ResetBudgetJob: ) ) verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e) + return _NO_PROGRESS + else: + return outcome @staticmethod async def _reset_expired_window( @@ -882,33 +1048,6 @@ class ResetBudgetJob: ) return user - @staticmethod - async def _reset_budget_for_enduser( - enduser: LiteLLM_EndUserTable, - ) -> LiteLLM_EndUserTable | None: - try: - enduser.spend = 0.0 - except Exception as e: - verbose_proxy_logger.exception("Error resetting budget for enduser: %s. Item: %s", e, enduser) - raise e - return enduser - - @staticmethod - async def _reset_budget_reset_at_date( - budget: LiteLLM_BudgetTableFull, - current_time: datetime, - reset_settings: BudgetResetSettings, - ) -> LiteLLM_BudgetTableFull: - try: - if budget.budget_duration is not None: - budget.budget_reset_at = compute_budget_reset_at( - budget_duration=budget.budget_duration, settings=reset_settings - ) - except Exception as e: - verbose_proxy_logger.exception("Error resetting budget_reset_at for budget: %s. Item: %s", e, budget) - raise e - return budget - @staticmethod async def _reset_budget_for_key( key: LiteLLM_VerificationToken, diff --git a/litellm/proxy/db/daily_spend_bulk_upsert.py b/litellm/proxy/db/daily_spend_bulk_upsert.py new file mode 100644 index 00000000000..55d325177c6 --- /dev/null +++ b/litellm/proxy/db/daily_spend_bulk_upsert.py @@ -0,0 +1,185 @@ +"""One multi-row ``INSERT ... ON CONFLICT DO UPDATE`` per batch of daily spend rows. + +Emitting a statement per aggregated key put every replica's flush on the database as +hundreds of separate statements against the same handful of hot rows, each holding its +row locks for the rest of the enclosing batch transaction. Folding a batch into a single +statement keeps the aggregation identical while collapsing both the statement count and +the window in which those locks are held. +""" + +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import groupby +from types import MappingProxyType +from typing import Final, Literal + +DailySpendEntity = Literal["user", "team", "org", "tag", "end_user", "agent"] + +SqlValue = str | int | float | None + +# A queued daily spend transaction, read by column name because the columns are data +# here rather than literals. The concrete TypedDicts in _types.py all satisfy this. +SpendRow = Mapping[str, object] + + +@dataclass(frozen=True, slots=True) +class DailySpendTable: + """The physical table behind one entity's daily rollup.""" + + name: str + entity_id_column: str + carries_request_id: bool = False + + +DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingProxyType( + { + "user": DailySpendTable(name="LiteLLM_DailyUserSpend", entity_id_column="user_id"), + "team": DailySpendTable(name="LiteLLM_DailyTeamSpend", entity_id_column="team_id"), + "org": DailySpendTable(name="LiteLLM_DailyOrganizationSpend", entity_id_column="organization_id"), + "end_user": DailySpendTable(name="LiteLLM_DailyEndUserSpend", entity_id_column="end_user_id"), + "agent": DailySpendTable(name="LiteLLM_DailyAgentSpend", entity_id_column="agent_id"), + "tag": DailySpendTable(name="LiteLLM_DailyTagSpend", entity_id_column="tag", carries_request_id=True), + } +) + +# The unique constraint's columns after the entity id, in constraint order. A NULL can +# never match itself in a unique index, so every one of these is normalized to '': the +# conflict target has to be NULL-free or the row is re-inserted on every single flush. +_KEY_COLUMNS: Final = ("date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") + +_COUNTER_COLUMNS: Final = ( + "prompt_tokens", + "completion_tokens", + "api_requests", + "successful_requests", + "failed_requests", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "compression_saved_tokens", +) +_SPEND_COLUMNS: Final = ( + "spend", + "compression_savings_spend", + "prompt_caching_savings_spend", + "autorouter_savings_spend", +) + +_CASTS: Final[Mapping[str, str]] = MappingProxyType( + { + **{column: "bigint" for column in _COUNTER_COLUMNS}, + **{column: "double precision" for column in _SPEND_COLUMNS}, + } +) + + +def _quoted(columns: Sequence[str]) -> str: + return ", ".join(f'"{column}"' for column in columns) + + +def _as_text(value: object) -> str: + return "" if value is None else str(value) + + +def _as_int(value: object) -> int: + return int(value) if isinstance(value, (int, float)) else 0 + + +def _as_float(value: object) -> float: + return float(value) if isinstance(value, (int, float)) else 0.0 + + +def conflict_key(table: DailySpendTable, transaction: SpendRow) -> tuple[str, ...]: + """The tuple the database arbitrates the upsert on, normalized free of NULLs.""" + return tuple(_as_text(transaction.get(column)) for column in (table.entity_id_column, *_KEY_COLUMNS)) + + +def _merge(group: Sequence[SpendRow]) -> SpendRow: + if len(group) == 1: + return group[0] + return { + **group[0], + **{column: sum(_as_int(row.get(column)) for row in group) for column in _COUNTER_COLUMNS}, + **{column: sum(_as_float(row.get(column)) for row in group) for column in _SPEND_COLUMNS}, + } + + +def merge_by_conflict_key( + table: DailySpendTable, + transactions: Sequence[SpendRow], +) -> tuple[tuple[tuple[str, ...], SpendRow], ...]: + """Batch entries keyed by the conflict tuple, in a deterministic order. + + The queue keys transactions by their raw field values, so two entries differing only + in a NULL versus an empty member reach the writer separately while arbitrating to the + same row. Postgres rejects a statement whose values touch one row twice, so they are + summed here into the single row they were always destined to become. Ordering by the + key keeps concurrent writers taking row locks in the same sequence. + """ + ordered: Final = sorted(transactions, key=lambda transaction: conflict_key(table, transaction)) + return tuple((key, _merge(tuple(group))) for key, group in groupby(ordered, key=lambda t: conflict_key(table, t))) + + +def _row_params( + table: DailySpendTable, + key: tuple[str, ...], + transaction: SpendRow, +) -> tuple[SqlValue, ...]: + request_id: Final = transaction.get("request_id") + return ( + str(uuid.uuid4()), + *key, + None if transaction.get("model_group") is None else _as_text(transaction.get("model_group")), + *(_as_int(transaction.get(column)) for column in _COUNTER_COLUMNS), + *(_as_float(transaction.get(column)) for column in _SPEND_COLUMNS), + *((None if request_id is None else _as_text(request_id),) if table.carries_request_id else ()), + ) + + +def _insert_columns(table: DailySpendTable) -> tuple[str, ...]: + return ( + "id", + table.entity_id_column, + *_KEY_COLUMNS, + "model_group", + *_COUNTER_COLUMNS, + *_SPEND_COLUMNS, + *(("request_id",) if table.carries_request_id else ()), + ) + + +def build_bulk_upsert( + table: DailySpendTable, + batch: Sequence[tuple[tuple[str, ...], SpendRow]], +) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one merged batch, plus its positional arguments.""" + columns: Final = _insert_columns(table) + quoted_table: Final = f'"{table.name}"' + rows: Final = ", ".join( + "(" + + ", ".join( + f"${row_index * len(columns) + offset + 1}::{_CASTS.get(column, 'text')}" + for offset, column in enumerate(columns) + ) + + ", (NOW() AT TIME ZONE 'UTC'))" + for row_index in range(len(batch)) + ) + increments: Final = ", ".join( + f'"{column}" = {quoted_table}."{column}" + EXCLUDED."{column}"' + for column in (*_COUNTER_COLUMNS, *_SPEND_COLUMNS) + ) + # request_id names one arbitrary contributing request, so an entry carrying none must + # not blank out the one already recorded. + request_id_update: Final = ( + f', "request_id" = COALESCE(EXCLUDED."request_id", {quoted_table}."request_id")' + if table.carries_request_id + else "" + ) + sql: Final = ( + f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n' + f"VALUES {rows}\n" + f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n" + f" {increments}{request_id_update},\n" + f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')" + ) + return sql, tuple(value for key, transaction in batch for value in _row_params(table, key, transaction)) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 385a21976b7..b0130db232a 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,9 +12,7 @@ import os import random import time import traceback -from collections.abc import Mapping from datetime import datetime, timedelta, timezone -from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload import litellm @@ -41,6 +39,11 @@ from litellm.proxy._types import ( SpendUpdateQueueItem, ToolDiscoveryQueueItem, ) +from litellm.proxy.db.daily_spend_bulk_upsert import ( + DAILY_SPEND_TABLES, + build_bulk_upsert, + merge_by_conflict_key, +) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) @@ -68,12 +71,6 @@ else: ProxyLogging = Any -# Only tag rows carry a request_id, so the other entity types spread nothing. Built -# once here rather than as an empty literal per transaction, and read-only so it cannot -# be filled in by accident from one of the call sites that spreads it. -_NO_TAG_REQUEST_ID: Final[Mapping[str, Any]] = MappingProxyType({}) - - def _get_llm_router(): """The proxy's router, or None outside a running proxy. @@ -1437,8 +1434,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyUserSpendTransaction], entity_type: Literal["user"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1451,8 +1446,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyTeamSpendTransaction], entity_type: Literal["team"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1465,8 +1458,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyOrganizationSpendTransaction], entity_type: Literal["org"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1479,8 +1470,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyEndUserSpendTransaction], entity_type: Literal["end_user"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1493,8 +1482,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyAgentSpendTransaction], entity_type: Literal["agent"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... @@ -1507,8 +1494,6 @@ class DBSpendUpdateWriter: daily_spend_transactions: dict[str, DailyTagSpendTransaction], entity_type: Literal["tag"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: ... # fmt: on @@ -1526,8 +1511,6 @@ class DBSpendUpdateWriter: | dict[str, DailyAgentSpendTransaction], entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"], entity_id_field: str, - table_name: str, - unique_constraint_name: str, ) -> None: """ Generic function to update daily spend for any entity type (user, team, org, tag, end_user, agent) @@ -1573,111 +1556,23 @@ class DBSpendUpdateWriter: ) return + table = DAILY_SPEND_TABLES[entity_type] try: - async with prisma_client.db.batch_() as batcher: - for _, transaction in transactions_to_process.items(): - entity_id = transaction.get(entity_id_field) - - # Construct the where clause dynamically - where_clause = { - unique_constraint_name: { - entity_id_field: entity_id, - "date": transaction["date"], - "api_key": transaction["api_key"], - "model": transaction["model"], - "custom_llm_provider": transaction.get("custom_llm_provider") or "", - "mcp_namespaced_tool_name": transaction.get("mcp_namespaced_tool_name") - or "", - "endpoint": transaction.get("endpoint") or "", - } - } - - # Get the table dynamically - table = getattr(batcher, table_name) - - # Additive metrics that older queued rows may omit; one - # enumeration feeds both the create and the increment below - optional_metrics = { - field: value - for field, value in ( - ("cache_read_input_tokens", transaction.get("cache_read_input_tokens")), - ( - "cache_creation_input_tokens", - transaction.get("cache_creation_input_tokens"), - ), - ("compression_saved_tokens", transaction.get("compression_saved_tokens")), - ( - "compression_savings_spend", - transaction.get("compression_savings_spend"), - ), - ( - "prompt_caching_savings_spend", - transaction.get("prompt_caching_savings_spend"), - ), - ("autorouter_savings_spend", transaction.get("autorouter_savings_spend")), - ) - if value is not None - } - - # Only tag rows carry a request_id. Resolved to a spreadable - # value here so both payloads are built in one shot: a dict - # appended to after construction is one nobody can reason about - # by reading its literal. - tag_request_id: Mapping[str, Any] = ( - MappingProxyType({"request_id": transaction["request_id"]}) - if entity_type == "tag" and "request_id" in transaction - else _NO_TAG_REQUEST_ID - ) - - # Common data structure for both create and update - common_data = { - entity_id_field: entity_id, - "date": transaction["date"], - "api_key": transaction["api_key"], - "model": transaction.get("model"), - "model_group": transaction.get("model_group"), - "mcp_namespaced_tool_name": transaction.get("mcp_namespaced_tool_name") or "", - "custom_llm_provider": transaction.get("custom_llm_provider"), - "endpoint": transaction.get("endpoint") or "", - "prompt_tokens": transaction["prompt_tokens"], - "completion_tokens": transaction["completion_tokens"], - "spend": transaction["spend"], - "api_requests": transaction["api_requests"], - "successful_requests": transaction["successful_requests"], - "failed_requests": transaction["failed_requests"], - **optional_metrics, - **tag_request_id, - } - - update_data = { - "prompt_tokens": {"increment": transaction["prompt_tokens"]}, - "completion_tokens": {"increment": transaction["completion_tokens"]}, - "spend": {"increment": transaction["spend"]}, - "api_requests": {"increment": transaction["api_requests"]}, - "successful_requests": {"increment": transaction["successful_requests"]}, - "failed_requests": {"increment": transaction["failed_requests"]}, - **{field: {"increment": value} for field, value in optional_metrics.items()}, - # An existing row predating the endpoint column gets it filled in here - "endpoint": transaction.get("endpoint") or "", - **tag_request_id, - } - - table.upsert( - where=where_clause, - data={ - "create": common_data, - "update": update_data, - }, - ) + # One statement per batch rather than per key: the same rows are + # aggregated, but concurrent writers no longer hold a batch's worth + # of row locks across a hundred round trips. + merged_batch = merge_by_conflict_key( + table=table, transactions=tuple(transactions_to_process.values()) + ) + sql, params = build_bulk_upsert(table=table, batch=merged_batch) + await prisma_client.db.execute_raw(sql, *params) except Exception as batch_error: # Log detailed error information for debugging batch upsert failures # This helps diagnose issues like unique constraint violations spend_log_error( - "Daily %s spend batch upsert failed. " - "Table: %s, Constraint: %s, Batch size: %d, Error: %s", + "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", entity_type, - table_name, - unique_constraint_name, + table.name, len(transactions_to_process), str(batch_error), exc=batch_error, @@ -1733,8 +1628,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1754,8 +1647,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="team", entity_id_field="team_id", - table_name="litellm_dailyteamspend", - unique_constraint_name="team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1775,8 +1666,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="org", entity_id_field="organization_id", - table_name="litellm_dailyorganizationspend", - unique_constraint_name="organization_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1796,8 +1685,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="end_user", entity_id_field="end_user_id", - table_name="litellm_dailyenduserspend", - unique_constraint_name="end_user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1817,8 +1704,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="agent", entity_id_field="agent_id", - table_name="litellm_dailyagentspend", - unique_constraint_name="agent_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @staticmethod @@ -1838,8 +1723,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_spend_transactions, entity_type="tag", entity_id_field="tag", - table_name="litellm_dailytagspend", - unique_constraint_name="tag_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) async def _common_add_spend_log_transaction_to_daily_transaction( diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 8991c3e0125..e0a21ceed26 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -1,5 +1,5 @@ from collections.abc import Awaitable, Callable -from typing import Any, Final +from typing import Any, Final, TypeVar from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -311,14 +311,17 @@ def _coerce_timeout(value: Any, fallback: float) -> float: return fallback +_ReadResultT: Final = TypeVar("_ReadResultT") + + async def call_with_db_reconnect_retry( prisma_client: Any, - coro_factory: Callable[[], Awaitable[Any]], + coro_factory: Callable[[], Awaitable[_ReadResultT]], *, reason: str, timeout_seconds: float | None = None, lock_timeout_seconds: float | None = None, -) -> Any: +) -> _ReadResultT: """Run a Prisma read coroutine with one transport-reconnect-and-retry. The canonical "self-heal a transient DB transport blip" wrapper used by diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 9c6dd32f15a..722f96ef814 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -9,10 +9,10 @@ import asyncio import json import os import re -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Coroutine, Mapping, Sequence from datetime import datetime from re import Pattern -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, cast import yaml from fastapi import HTTPException @@ -28,6 +28,7 @@ from litellm.types.utils import ( GenericGuardrailAPIInputs, GuardrailStatus, GuardrailTracingDetail, + ModelResponse, ModelResponseStream, ) @@ -83,6 +84,46 @@ WORD_NUMBER_SEQUENCE_PATTERN: Final = re.compile( WORD_NUMBER_TOKEN_FINDER: Final = re.compile(rf"(?:{WORD_NUMBER_TOKEN_REGEX})", re.IGNORECASE) +class ConditionalCategoryConfig(TypedDict): + identifier_words: Sequence[str] + block_words: Sequence[str] + action: ContentFilterAction + severity: str + + +class CompiledPatternEntry(TypedDict): + regex: Pattern[str] + pattern_name: str + action: ContentFilterAction + keyword_regex: Pattern[str] | None + allow_word_numbers: bool + + +class _PatternExtraLookup(TypedDict): + keyword_pattern: str | None + allow_word_numbers: bool + + +class _CategoryConfigView(TypedDict): + category: object + enabled: object + action: object + category_file: str | None + + +class CategoryFileData(TypedDict, total=False): + category_name: str + description: str + default_action: str + keywords: Sequence[Mapping[str, str]] + exceptions: Sequence[str] + identifier_words: Sequence[str] + always_block_keywords: Sequence[Mapping[str, str]] + inherit_from: str + additional_block_words: Sequence[str] + phrase_patterns: Sequence[str] + + # Helper data structure for category-based detection class CategoryConfig: """Configuration for a content category.""" @@ -92,13 +133,13 @@ class CategoryConfig: category_name: str, description: str, default_action: ContentFilterAction, - keywords: list[dict[str, str]], - exceptions: list[str], - identifier_words: list[str] | None = None, - always_block_keywords: list[dict[str, str]] | None = None, + keywords: Sequence[Mapping[str, str]], + exceptions: Sequence[str], + identifier_words: Sequence[str] | None = None, + always_block_keywords: Sequence[Mapping[str, str]] | None = None, inherit_from: str | None = None, - additional_block_words: list[str] | None = None, - phrase_patterns: list[str] | None = None, + additional_block_words: Sequence[str] | None = None, + phrase_patterns: Sequence[str] | None = None, ): self.category_name = category_name self.description = description @@ -151,7 +192,7 @@ class ContentFilterGuardrail(CustomGuardrail): severity_threshold: str = "medium", llm_router: Router | None = None, image_model: str | None = None, - competitor_intent_config: dict[str, Any] | None = None, + competitor_intent_config: dict[str, object] | None = None, **kwargs, ): """ @@ -194,9 +235,7 @@ class ContentFilterGuardrail(CustomGuardrail): # Always-block keywords are checked after exceptions (exceptions take precedence) self.always_block_category_keywords: dict[str, tuple[str, str, ContentFilterAction]] = {} # Store conditional categories (identifier_words + block_words) - self.conditional_categories: dict[ - str, dict[str, Any] - ] = {} # category_name -> {identifier_words, block_words, action, severity} + self.conditional_categories: dict[str, ConditionalCategoryConfig] = {} # Competitor intent checker (optional; airline uses major_airlines.json, generic requires competitors) self._competitor_intent_checker: BaseCompetitorIntentChecker | None = None @@ -212,7 +251,7 @@ class ContentFilterGuardrail(CustomGuardrail): normalized_blocked_words: Final = self._normalize_blocked_words(blocked_words) # Compile regex patterns - self.compiled_patterns: list[dict[str, Any]] = [] + self.compiled_patterns: list[CompiledPatternEntry] = [] for pattern_config in normalized_patterns: self._add_pattern(pattern_config) @@ -250,7 +289,7 @@ class ContentFilterGuardrail(CustomGuardrail): "Loaded %s categories with %s keywords", len(self.loaded_categories), len(self.category_keywords) ) - def _init_competitor_intent_checker(self, competitor_intent_config: dict[str, Any]) -> None: + def _init_competitor_intent_checker(self, competitor_intent_config: dict[str, object]) -> None: try: competitor_intent_type: Final = competitor_intent_config.get("competitor_intent_type", "airline") if competitor_intent_type == "generic": @@ -293,6 +332,15 @@ class ContentFilterGuardrail(CustomGuardrail): result.append(word) return result + @staticmethod + def _category_config_view(cat_config: ContentFilterCategoryConfig) -> _CategoryConfigView: + return { + "category": cat_config.get("category"), + "enabled": cat_config.get("enabled", True), + "action": cat_config.get("action"), + "category_file": cat_config.get("category_file"), + } + @staticmethod def _assert_within_categories_dir(path: str, categories_dir: str) -> None: """Raise ValueError if path escapes the categories directory.""" @@ -395,7 +443,8 @@ class ContentFilterGuardrail(CustomGuardrail): categories_dir: Final = os.path.join(os.path.dirname(__file__), "categories") for cat_config in categories: - category_name = cat_config.get("category") + view = self._category_config_view(cat_config) + category_name = view["category"] if not category_name or not isinstance(category_name, str): verbose_proxy_logger.warning("Category name missing or invalid in config, skipping") continue @@ -405,12 +454,12 @@ class ContentFilterGuardrail(CustomGuardrail): verbose_proxy_logger.warning("Category name '%s' contains invalid characters, skipping", category_name) continue - enabled = cat_config.get("enabled", True) - action = cat_config.get("action") + enabled = view["enabled"] + action = view["action"] severity_threshold = ( cat_config.get("severity_threshold", self.severity_threshold) or self.severity_threshold ) - custom_file = cat_config.get("category_file") + custom_file = view["category_file"] if not enabled: verbose_proxy_logger.debug("Category %s is disabled, skipping", category_name) @@ -514,7 +563,7 @@ class ContentFilterGuardrail(CustomGuardrail): categories_dir: Directory containing category files """ try: - block_words: Final = [] + block_words: Final[list[str]] = [] inherit_from = category_config_obj.inherit_from # Load inherited block words if specified @@ -605,11 +654,7 @@ class ContentFilterGuardrail(CustomGuardrail): """ if file_path.lower().endswith(".json"): return self._load_category_file_json(file_path) - with open(file_path, "r") as f: - data: Final = yaml.safe_load(f) - - # Handle always_block_keywords if present - always_block: Final = data.get("always_block_keywords", []) + data: Final = self._read_category_yaml(file_path) return CategoryConfig( category_name=data.get("category_name", "unknown"), @@ -618,12 +663,17 @@ class ContentFilterGuardrail(CustomGuardrail): keywords=data.get("keywords", []), exceptions=data.get("exceptions", []), identifier_words=data.get("identifier_words"), - always_block_keywords=always_block, + always_block_keywords=data.get("always_block_keywords", []), inherit_from=data.get("inherit_from"), additional_block_words=data.get("additional_block_words"), phrase_patterns=data.get("phrase_patterns"), ) + @staticmethod + def _read_category_yaml(file_path: str) -> CategoryFileData: + with open(file_path, "r") as f: + return yaml.safe_load(f) + def _load_category_file_json(self, file_path: str) -> CategoryConfig: """ Load a category from the harm_toxic_abuse-style JSON format. @@ -682,13 +732,13 @@ class ContentFilterGuardrail(CustomGuardrail): pattern_config: ContentFilterPattern configuration """ try: - extra_config: dict[str, Any] = {} + extra_config: _PatternExtraLookup = {"keyword_pattern": None, "allow_word_numbers": False} if pattern_config.pattern_type == "prebuilt": if not pattern_config.pattern_name: raise ValueError("pattern_name is required for prebuilt patterns") compiled = get_compiled_pattern(pattern_config.pattern_name) pattern_name = pattern_config.pattern_name - extra_config = PATTERN_EXTRA_CONFIG.get(pattern_name, {}) or {} + extra_config = self._lookup_pattern_extra(pattern_name) elif pattern_config.pattern_type == "regex": if not pattern_config.pattern: raise ValueError("pattern is required for regex patterns") @@ -697,9 +747,8 @@ class ContentFilterGuardrail(CustomGuardrail): else: raise ValueError(f"Unknown pattern_type: {pattern_config.pattern_type}") - keyword_regex: Pattern | None = None - if extra_config.get("keyword_pattern"): - keyword_regex = re.compile(extra_config["keyword_pattern"], re.IGNORECASE) + keyword_pattern: Final = extra_config["keyword_pattern"] + keyword_regex: Final = re.compile(keyword_pattern, re.IGNORECASE) if keyword_pattern else None self.compiled_patterns.append( { @@ -707,7 +756,7 @@ class ContentFilterGuardrail(CustomGuardrail): "pattern_name": pattern_name, "action": pattern_config.action, "keyword_regex": keyword_regex, - "allow_word_numbers": bool(extra_config.get("allow_word_numbers")), + "allow_word_numbers": extra_config["allow_word_numbers"], } ) verbose_proxy_logger.debug("Added pattern: %s with action %s", pattern_name, pattern_config.action) @@ -715,6 +764,14 @@ class ContentFilterGuardrail(CustomGuardrail): verbose_proxy_logger.error("Error adding pattern %s: %s", pattern_config, e) raise + @staticmethod + def _lookup_pattern_extra(pattern_name: str) -> _PatternExtraLookup: + extra: Final = PATTERN_EXTRA_CONFIG.get(pattern_name) + return { + "keyword_pattern": extra.get("keyword_pattern") if extra is not None else None, + "allow_word_numbers": bool(extra.get("allow_word_numbers")) if extra is not None else False, + } + def _load_blocked_words_file(self, file_path: str) -> None: """ Load blocked words from a YAML file. @@ -754,18 +811,16 @@ class ContentFilterGuardrail(CustomGuardrail): except Exception as e: raise Exception(f"Error loading blocked words file {file_path}: {e}") - def _find_pattern_spans(self, text: str, pattern_entry: dict[str, Any]) -> list[tuple[int, int]]: + def _find_pattern_spans(self, text: str, pattern_entry: CompiledPatternEntry) -> list[tuple[int, int]]: """Return all match spans for a pattern, applying contextual rules if required.""" - regex: Final[Pattern] = pattern_entry["regex"] - keyword_regex: Final[Pattern | None] = pattern_entry.get("keyword_regex") + regex: Final[Pattern[str]] = pattern_entry["regex"] + keyword_regex: Final[Pattern[str] | None] = pattern_entry.get("keyword_regex") allow_word_numbers: Final[bool] = pattern_entry.get("allow_word_numbers", False) - keyword_matches: list[re.Match] | None = None - if keyword_regex is not None: - keyword_matches = list(keyword_regex.finditer(text)) - if not keyword_matches: - return [] + keyword_matches: Final = list(keyword_regex.finditer(text)) if keyword_regex is not None else None + if keyword_matches is not None and not keyword_matches: + return [] match_spans: Final[list[tuple[int, int]]] = [] @@ -795,7 +850,7 @@ class ContentFilterGuardrail(CustomGuardrail): self, value_start: int, value_end: int, - keyword_matches: list[re.Match], + keyword_matches: Sequence[re.Match[str]], text: str, ) -> bool: """Check if a value is separated from a keyword by an allowed gap.""" @@ -861,7 +916,7 @@ class ContentFilterGuardrail(CustomGuardrail): def _convert_word_number_sequence(self, sequence: str) -> str | None: """Convert a spelled-out digit sequence (e.g., 'One-Two') into digits.""" - tokens: Final = WORD_NUMBER_TOKEN_FINDER.findall(sequence) + tokens: Final[list[str]] = WORD_NUMBER_TOKEN_FINDER.findall(sequence) if not tokens: return None @@ -1328,7 +1383,7 @@ class ContentFilterGuardrail(CustomGuardrail): HTTPException: If sensitive content is detected and action is BLOCK """ # Collect all exceptions from loaded categories - all_exceptions: Final = [] + all_exceptions: Final[list[str]] = [] for category in self.loaded_categories.values(): all_exceptions.extend(category.exceptions) @@ -1404,7 +1459,7 @@ class ContentFilterGuardrail(CustomGuardrail): if not (images and self.image_model and self.llm_router): return - tasks: Final = [] + tasks: Final[list[Coroutine[object, object, ModelResponse]]] = [] for image in images: task = self.llm_router.acompletion( model=self.image_model, @@ -1425,12 +1480,10 @@ class ContentFilterGuardrail(CustomGuardrail): tasks.append(task) responses: Final = await asyncio.gather(*tasks) - descriptions: Final = [] + descriptions: Final[list[str]] = [] for response in responses: - choice = response.choices[0] - message = getattr(choice, "message", None) - if message and getattr(message, "content", None): - image_description = message.content + image_description = self._describe_image_response_content(response) + if image_description: verbose_proxy_logger.debug("Image description: %s", image_description) descriptions.append(image_description) else: @@ -1447,7 +1500,7 @@ class ContentFilterGuardrail(CustomGuardrail): except HTTPException as e: # e.detail can be a string or dict if isinstance(e.detail, dict) and "error" in e.detail: - detail_dict = cast(dict[str, Any], e.detail) + detail_dict = cast(dict[str, str], e.detail) detail_dict["error"] = detail_dict["error"] + " (Image description): " + description elif isinstance(e.detail, str): e.detail = e.detail + " (Image description): " + description @@ -1455,6 +1508,14 @@ class ContentFilterGuardrail(CustomGuardrail): e.detail = "Content blocked: Image description detected" + description raise e + @staticmethod + def _describe_image_response_content(response: ModelResponse) -> str | None: + choice = response.choices[0] + message = getattr(choice, "message", None) + if message and getattr(message, "content", None): + return message.content + return None + def _count_masked_entities( self, detections: list[ContentFilterDetection], @@ -1484,12 +1545,12 @@ class ContentFilterGuardrail(CustomGuardrail): category = category_detection["category"] masked_entity_count[category] = masked_entity_count.get(category, 0) + 1 - def _build_match_details(self, detections: list[ContentFilterDetection]) -> list[dict]: + def _build_match_details(self, detections: list[ContentFilterDetection]) -> list[dict[str, object]]: """Build match_details list from content filter detections.""" - match_details: Final[list[dict]] = [] + match_details: Final[list[dict[str, object]]] = [] for detection in detections: action_taken = detection.get("action", detection.get("action_hint", "")) - detail: dict = {"type": detection["type"], "action_taken": action_taken} + detail: dict[str, object] = {"type": detection["type"], "action_taken": action_taken} if detection["type"] == "pattern": detail["detection_method"] = "regex" detail["snippet"] = cast(PatternDetection, detection).get("pattern_name", "") @@ -1510,7 +1571,7 @@ class ContentFilterGuardrail(CustomGuardrail): def _get_detection_methods(self, detections: list[ContentFilterDetection]) -> str: """Get comma-separated detection methods used.""" - methods: Final[set] = set() + methods: Final[set[str]] = set() for detection in detections: if detection["type"] == "pattern": methods.add("regex") @@ -1659,7 +1720,7 @@ class ContentFilterGuardrail(CustomGuardrail): guardrail_json_response = exception_str if exception_str else [dict(detection) for detection in detections] # Competitor intent: add confidence and classification to tracing if present - tracing_kw: Final[dict[str, Any]] = { + tracing_kw: Final[GuardrailTracingDetail] = { "guardrail_id": self.config_guardrail_id or self.guardrail_name, "policy_template": self.config_policy_template or self._get_policy_templates(), "detection_method": (self._get_detection_methods(detections) if detections else None), diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py index be6678862ec..6c23813affd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/patterns.py @@ -55,7 +55,7 @@ for pattern_data in _PATTERNS_DATA["patterns"]: PATTERN_EXTRA_CONFIG[pattern_data["name"]] = extra_config -def get_compiled_pattern(pattern_name: str) -> Pattern: +def get_compiled_pattern(pattern_name: str) -> Pattern[str]: """ Get a compiled regex pattern by name. diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py index 01ca785ad68..e2d7c06f7c5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_jwt_signer/mcp_jwt_signer.py @@ -73,12 +73,14 @@ import hashlib import os import re import time -from typing import Any, Final, Optional +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Optional import jwt from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa -from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey +from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey +from typing_extensions import NotRequired, TypedDict from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -90,13 +92,28 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral +if TYPE_CHECKING: + from jwt.types import Options + + +class _OIDCDiscoveryDocument(TypedDict, total=False): + jwks_uri: str + + +class _JWTDecodeKwargs(TypedDict): + algorithms: Sequence[str] + options: "Options" + audience: NotRequired[str] + issuer: NotRequired[str] + + # Module-level singleton for the JWKS discovery endpoint to access. _mcp_jwt_signer_instance: Optional["MCPJWTSigner"] = None _MCP_JWT_CALL_TYPES: Final = frozenset({"call_mcp_tool", "list_mcp_tools"}) # Simple in-memory JWKS cache: keyed by JWKS URI → (keys_list, fetched_at). -_jwks_cache: Final[dict[str, tuple]] = {} +_jwks_cache: Final[dict[str, tuple[Sequence[Mapping[str, object]], float]]] = {} _JWKS_CACHE_TTL: Final = 3600 # 1 hour @@ -133,7 +150,7 @@ def _int_to_base64url(n: int) -> str: return base64.urlsafe_b64encode(n.to_bytes(byte_length, byteorder="big")).rstrip(b"=").decode("ascii") -def _compute_kid(public_key: Any) -> str: +def _compute_kid(public_key: RSAPublicKey) -> str: """Derive a key ID from the public key's DER encoding (SHA-256, first 16 hex chars).""" der_bytes: Final = public_key.public_bytes( encoding=serialization.Encoding.DER, @@ -142,7 +159,7 @@ def _compute_kid(public_key: Any) -> str: return hashlib.sha256(der_bytes).hexdigest()[:16] -async def _fetch_jwks(jwks_uri: str) -> list[dict[str, Any]]: +async def _fetch_jwks(jwks_uri: str) -> Sequence[Mapping[str, object]]: """ Fetch and cache a JWKS from the given URI. @@ -163,12 +180,13 @@ async def _fetch_jwks(jwks_uri: str) -> list[dict[str, Any]]: client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) resp: Final = await client.get(jwks_uri, headers={"Accept": "application/json"}) resp.raise_for_status() - keys = resp.json().get("keys", []) - _jwks_cache[jwks_uri] = (keys, now) - return keys + jwks_body: Final[Mapping[str, Sequence[Mapping[str, object]]]] = resp.json() + fetched_keys: Final = jwks_body.get("keys", []) + _jwks_cache[jwks_uri] = (fetched_keys, now) + return fetched_keys -async def _fetch_oidc_discovery(discovery_uri: str) -> dict[str, Any]: +async def _fetch_oidc_discovery(discovery_uri: str) -> _OIDCDiscoveryDocument: """Fetch an OIDC discovery document and return its parsed JSON.""" from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -178,7 +196,8 @@ async def _fetch_oidc_discovery(discovery_uri: str) -> dict[str, Any]: client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check) resp: Final = await client.get(discovery_uri, headers={"Accept": "application/json"}) resp.raise_for_status() - return resp.json() + document: Final[_OIDCDiscoveryDocument] = resp.json() + return document class MCPJWTSigner(CustomGuardrail): @@ -230,8 +249,8 @@ class MCPJWTSigner(CustomGuardrail): # FR-12: End-user identity mapping end_user_claim_sources: list[str] | None = None, # FR-13: Claim operations - add_claims: dict[str, Any] | None = None, - set_claims: dict[str, Any] | None = None, + add_claims: Mapping[str, object] | None = None, + set_claims: Mapping[str, object] | None = None, remove_claims: list[str] | None = None, # FR-14: Two-token model channel_token_audience: str | None = None, @@ -283,7 +302,7 @@ class MCPJWTSigner(CustomGuardrail): self.verify_issuer: str | None = verify_issuer self.verify_audience: str | None = verify_audience # Cached OIDC discovery document (fetched lazily, TTL = 24 h) - self._oidc_discovery_doc: dict[str, Any] | None = None + self._oidc_discovery_doc: _OIDCDiscoveryDocument | None = None self._oidc_discovery_fetched_at: float = 0.0 # --- FR-12: End-user identity mapping --- @@ -294,8 +313,8 @@ class MCPJWTSigner(CustomGuardrail): ] # --- FR-13: Claim operations --- - self.add_claims: dict[str, Any] = add_claims or {} - self.set_claims: dict[str, Any] = set_claims or {} + self.add_claims: Mapping[str, object] = add_claims or {} + self.set_claims: Mapping[str, object] = set_claims or {} self.remove_claims: list[str] = remove_claims or [] # --- FR-14: Two-token model --- @@ -347,7 +366,7 @@ class MCPJWTSigner(CustomGuardrail): """ return 3600 if self._persistent_key else 300 - def get_jwks(self) -> dict[str, Any]: + def get_jwks(self) -> Mapping[str, Sequence[Mapping[str, str]]]: """ Return the JWKS for the RSA public key. Used by GET /.well-known/jwks.json so MCP servers can verify tokens. @@ -374,7 +393,7 @@ class MCPJWTSigner(CustomGuardrail): # the IdP, short enough to pick up jwks_uri changes after key rotation. _OIDC_DISCOVERY_TTL = 86400 - async def _get_oidc_discovery(self) -> dict[str, Any]: + async def _get_oidc_discovery(self) -> _OIDCDiscoveryDocument: """Fetch and cache the OIDC discovery document with a 24-hour TTL. Only caches when the doc contains a 'jwks_uri' so that a transient or @@ -391,7 +410,7 @@ class MCPJWTSigner(CustomGuardrail): return doc return self._oidc_discovery_doc or {} - async def _verify_incoming_jwt(self, raw_token: str) -> dict[str, Any]: + async def _verify_incoming_jwt(self, raw_token: str) -> dict[str, object]: """ Verify an incoming Bearer JWT against the configured IdP's JWKS. @@ -438,8 +457,8 @@ class MCPJWTSigner(CustomGuardrail): # it infers from the key type (RSAPublicKey → RS256). alg: Final = getattr(signing_jwk, "algorithm_name", None) or "RS256" - decode_options: Final[dict[str, Any]] = {"verify_exp": True} - decode_kwargs: Final[dict[str, Any]] = { + decode_options: Final[Options] = {"verify_exp": True} + decode_kwargs: Final[_JWTDecodeKwargs] = { "algorithms": [alg], "options": decode_options, } @@ -451,10 +470,10 @@ class MCPJWTSigner(CustomGuardrail): if self.verify_issuer: decode_kwargs["issuer"] = self.verify_issuer - payload: Final[dict[str, Any]] = jwt.decode(raw_token, signing_jwk.key, **decode_kwargs) + payload: Final[dict[str, object]] = jwt.decode(raw_token, signing_jwk.key, **decode_kwargs) return payload - async def _introspect_opaque_token(self, token: str) -> dict[str, Any]: + async def _introspect_opaque_token(self, token: str) -> dict[str, object]: """ Perform RFC 7662 token introspection for opaque (non-JWT) tokens. @@ -479,7 +498,7 @@ class MCPJWTSigner(CustomGuardrail): headers={"Accept": "application/json"}, ) resp.raise_for_status() - result: Final[dict[str, Any]] = resp.json() + result: Final[dict[str, object]] = resp.json() if not result.get("active", False): raise jwt.exceptions.ExpiredSignatureError( "MCPJWTSigner: incoming token is inactive (introspection returned active=false)" @@ -492,7 +511,7 @@ class MCPJWTSigner(CustomGuardrail): def _validate_required_claims( self, - jwt_claims: dict[str, Any] | None, + jwt_claims: Mapping[str, object] | None, ) -> None: """ Raise HTTP 403 if any required_claims are absent from the verified @@ -522,7 +541,7 @@ class MCPJWTSigner(CustomGuardrail): def _resolve_end_user_identity( self, user_api_key_dict: UserAPIKeyAuth, - jwt_claims: dict[str, Any] | None, + jwt_claims: Mapping[str, object] | None, ) -> str: """ Resolve the outbound JWT 'sub' using the ordered end_user_claim_sources list. @@ -545,19 +564,19 @@ class MCPJWTSigner(CustomGuardrail): value = str(raw) if raw else None elif source == "litellm:user_id": - uid = getattr(user_api_key_dict, "user_id", None) + uid = user_api_key_dict.user_id value = str(uid) if uid else None elif source == "litellm:email": - email = getattr(user_api_key_dict, "user_email", None) + email = user_api_key_dict.user_email value = str(email) if email else None elif source == "litellm:end_user_id": - eid = getattr(user_api_key_dict, "end_user_id", None) + eid = user_api_key_dict.end_user_id value = str(eid) if eid else None elif source == "litellm:team_id": - tid = getattr(user_api_key_dict, "team_id", None) + tid = user_api_key_dict.team_id value = str(tid) if tid else None else: @@ -568,7 +587,7 @@ class MCPJWTSigner(CustomGuardrail): return value # Final fallback for service accounts with no user identity - token: Final = getattr(user_api_key_dict, "token", None) or getattr(user_api_key_dict, "api_key", None) + token: Final = user_api_key_dict.token or user_api_key_dict.api_key if token: return "apikey:" + hashlib.sha256(str(token).encode()).hexdigest()[:16] return "litellm-proxy" @@ -615,7 +634,7 @@ class MCPJWTSigner(CustomGuardrail): # FR-13: Claim operations # ------------------------------------------------------------------ - def _apply_claim_operations(self, claims: dict[str, Any]) -> dict[str, Any]: + def _apply_claim_operations(self, claims: dict[str, object]) -> dict[str, object]: """Apply add_claims, set_claims, and remove_claims to the claim dict.""" # add_claims: insert only when key is absent for k, v in self.add_claims.items(): @@ -637,9 +656,9 @@ class MCPJWTSigner(CustomGuardrail): def _passthrough_optional_claims( self, - claims: dict[str, Any], - jwt_claims: dict[str, Any] | None, - ) -> dict[str, Any]: + claims: dict[str, object], + jwt_claims: Mapping[str, object] | None, + ) -> dict[str, object]: """Forward optional_claims from verified incoming token into the outbound JWT.""" if not self.optional_claims or not jwt_claims: return claims @@ -656,7 +675,7 @@ class MCPJWTSigner(CustomGuardrail): self, user_api_key_dict: UserAPIKeyAuth, data: dict, - jwt_claims: dict[str, Any] | None = None, + jwt_claims: Mapping[str, object] | None = None, call_type: CallTypesLiteral | None = None, ) -> dict[str, Any]: """ @@ -669,7 +688,7 @@ class MCPJWTSigner(CustomGuardrail): jwt_claims if available. None for pure API-key requests. """ now: Final = int(time.time()) - claims: dict[str, Any] = { + claims: dict[str, object] = { "iss": self.issuer, "aud": self.audience, "iat": now, @@ -681,18 +700,18 @@ class MCPJWTSigner(CustomGuardrail): claims["sub"] = self._resolve_end_user_identity(user_api_key_dict, jwt_claims) # email passthrough when available from LiteLLM context - user_email: Final = getattr(user_api_key_dict, "user_email", None) + user_email: Final = user_api_key_dict.user_email if user_email: claims["email"] = user_email # act — RFC 8693 delegation claim (team/org context) - team_id: Final = getattr(user_api_key_dict, "team_id", None) - org_id: Final = getattr(user_api_key_dict, "org_id", None) + team_id: Final = user_api_key_dict.team_id + org_id: Final = user_api_key_dict.org_id act_sub: Final = team_id or org_id or "litellm-proxy" claims["act"] = {"sub": act_sub} # end_user_id when set separately from user_id - end_user_id: Final = getattr(user_api_key_dict, "end_user_id", None) + end_user_id: Final = user_api_key_dict.end_user_id if end_user_id: claims["end_user_id"] = end_user_id @@ -710,8 +729,8 @@ class MCPJWTSigner(CustomGuardrail): def _build_channel_token_claims( self, - base_claims: dict[str, Any], - ) -> dict[str, Any]: + base_claims: Mapping[str, object], + ) -> dict[str, object]: """ Build claims for the channel token (FR-14 two-token model). @@ -776,7 +795,7 @@ class MCPJWTSigner(CustomGuardrail): # ------------------------------------------------------------------ # FR-5: Verify incoming token before re-signing # ------------------------------------------------------------------ - jwt_claims: dict[str, Any] | None = None + jwt_claims: dict[str, object] | None = None raw_token: Final[str | None] = hook_data.get("incoming_bearer_token") if self.access_token_discovery_uri and raw_token: @@ -810,7 +829,7 @@ class MCPJWTSigner(CustomGuardrail): # Fall back to LiteLLM-decoded JWT claims (available when proxy uses JWT auth). if jwt_claims is None: - jwt_claims = getattr(user_api_key_dict, "jwt_claims", None) + jwt_claims = user_api_key_dict.jwt_claims # ------------------------------------------------------------------ # FR-15: Validate required claims @@ -896,7 +915,7 @@ async def inject_mcp_jwt_headers_for_upstream( if auth_hdr.lower().startswith("bearer "): incoming_bearer_token = auth_hdr[len("bearer ") :] - hook_data: Final[dict[str, Any]] = { + hook_data: Final = { "mcp_tool_name": "" if for_list_tools else mcp_tool_name, "incoming_bearer_token": incoming_bearer_token, "extra_headers": merged, diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 13ced0ac06c..dcd86f98ee4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -8,6 +8,7 @@ Provides real-time threat detection, DLP, URL filtering, content masking, and po import json import os import re +from collections.abc import AsyncIterable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal, Optional from urllib.parse import urlparse @@ -166,7 +167,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): GuardrailEventHooks.during_mcp_call: GuardrailEventHooks.during_call, } - def should_run_guardrail(self, data: Any, event_type: GuardrailEventHooks) -> bool: + def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool: if super().should_run_guardrail(data, event_type): return True compat: Final = self._MCP_COMPAT_MAP.get(event_type) @@ -175,7 +176,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): return True return False - def _extract_text_from_messages(self, messages: list[dict[str, Any]]) -> str: + def _extract_text_from_messages(self, messages: Sequence[Mapping[str, object]]) -> str: """Extract text content from messages array.""" if not isinstance(messages, list) or not messages: return "" @@ -242,10 +243,10 @@ class PanwPrismaAirsHandler(CustomGuardrail): self, content: str = "", is_response: bool = False, - metadata: dict[str, Any] | None = None, - call_id: str | None = None, + metadata: Mapping[str, object] | None = None, + call_id: object = None, tool_event: dict[str, Any] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Call PANW Prisma AIRS API to scan content or a tool_event.""" if tool_event is None and not content.strip(): @@ -275,7 +276,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): else: app_name_value = self.app_name # Defaults to "LiteLLM" - panw_metadata: Final = { + panw_metadata: Final[dict[str, object]] = { "app_user": ( (metadata.get("app_user") or metadata.get("user") or "litellm_user") if metadata else "litellm_user" ), @@ -295,13 +296,13 @@ class PanwPrismaAirsHandler(CustomGuardrail): panw_metadata["litellm_trace_id"] = metadata["litellm_trace_id"] # Build contents: tool_event takes priority, else prompt/response text - contents: list[dict[str, Any]] + contents: Sequence[Mapping[str, object]] if tool_event is not None: contents = [{"tool_event": tool_event}] else: contents = [{"response" if is_response else "prompt": content}] - payload: Final = { + payload: Final[dict[str, object]] = { "metadata": panw_metadata, "contents": contents, } @@ -325,7 +326,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # If neither profile_name nor profile_id is provided, PANW API will use the # profile linked to the API key (if configured in Strata Cloud Manager) if profile_name or profile_id: - ai_profile: Final = {} + ai_profile: Final[dict[str, object]] = {} if profile_id: ai_profile["profile_id"] = profile_id if profile_name: @@ -333,7 +334,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): payload["ai_profile"] = ai_profile if is_response and tool_event is None: - payload["metadata"]["is_response"] = True + panw_metadata["is_response"] = True headers: Final = { "Content-Type": "application/json", @@ -355,7 +356,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) response.raise_for_status() - result: Final = response.json() + result: Final[dict[str, object]] = response.json() # Validate response format if "action" not in result: @@ -489,7 +490,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): ) return "unknown" - def _get_masked_text(self, scan_result: dict[str, Any], is_response: bool = False) -> str | None: + def _get_masked_text(self, scan_result: Mapping[str, object], is_response: bool = False) -> str | None: """Extract masked text from PANW scan result.""" masked_key: Final = "response_masked_data" if is_response else "prompt_masked_data" masked_data: Final = scan_result.get(masked_key) @@ -511,7 +512,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): @staticmethod def _apply_mcp_masking( request_data: dict, - original_args: Any, + original_args: object, masked_text: str, *, is_blocked: bool = True, @@ -544,7 +545,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # If the original args were structured, preserve the type. if isinstance(original_args, (dict, list)): try: - parsed: Final = json.loads(masked_text) + parsed: Final[object] = json.loads(masked_text) except (json.JSONDecodeError, TypeError): raise HTTPException( status_code=400, @@ -556,7 +557,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): } }, ) - masked_value: Any = parsed + masked_value: object = parsed else: masked_value = masked_text @@ -572,7 +573,9 @@ class PanwPrismaAirsHandler(CustomGuardrail): else: verbose_proxy_logger.info("PANW Prisma AIRS: MCP request allowed with PII masking applied") - def _apply_masking_to_messages(self, messages: list[dict[str, Any]], masked_text: str) -> list[dict[str, Any]]: + def _apply_masking_to_messages( + self, messages: list[dict[str, object]], masked_text: str + ) -> Sequence[Mapping[str, object]]: """Apply masked text to the last user message.""" if not messages: return messages @@ -622,7 +625,9 @@ class PanwPrismaAirsHandler(CustomGuardrail): if hasattr(choice.message.function_call, "arguments"): choice.message.function_call.arguments = masked_text - def _build_error_detail(self, scan_result: dict[str, Any], is_response: bool = False) -> dict[str, Any]: + def _build_error_detail( + self, scan_result: Mapping[str, object], is_response: bool = False + ) -> Mapping[str, Mapping[str, object]]: """Build enhanced error detail with scan information.""" action_type: Final = "Response" if is_response else "Prompt" code_suffix: Final = "_response_blocked" if is_response else "_blocked" @@ -642,7 +647,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): }, ) - error_detail: Final = { + error_detail: Final[dict[str, dict[str, object]]] = { "error": { "message": error_msg, "type": "guardrail_violation", @@ -672,12 +677,12 @@ class PanwPrismaAirsHandler(CustomGuardrail): def _handle_api_error_with_logging( self, - scan_result: dict[str, Any], - data: dict[str, Any], + scan_result: dict[str, object], + data: dict[str, object], start_time: datetime, event_type: GuardrailEventHooks, is_response: bool = False, - ) -> dict[str, Any] | None: + ) -> None: """Handle API errors with fail-open/fail-closed logic.""" end_time: Final = datetime.now() duration: Final = (end_time - start_time).total_seconds() @@ -722,7 +727,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=f"{self.guardrail_name}:unscanned" ) - return None + return raise HTTPException( status_code=500, @@ -783,7 +788,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): return metadata @staticmethod - def _extract_text_from_sse_bytes(chunks: list[bytes]) -> str: + def _extract_text_from_sse_bytes(chunks: Sequence[bytes]) -> str: """Extract text from Anthropic SSE byte chunks (content_block_delta → text_delta).""" texts: Final[list[str]] = [] raw: Final = b"".join(chunks).decode("utf-8", errors="replace") @@ -804,7 +809,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): return "".join(texts) @staticmethod - def _extract_text_from_streaming_events(chunks: list) -> str: + def _extract_text_from_streaming_events(chunks: Sequence[object]) -> str: """Extract text from /v1/responses streaming events (object or dict).""" def _attr(c, key): @@ -960,7 +965,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): cache: DualCache, data: dict[str, Any], call_type: CallTypesLiteral, - ) -> dict[str, Any] | None: + ) -> dict[str, object] | None: """ Pre-call hook to scan user prompts before sending to LLM. @@ -1075,10 +1080,10 @@ class PanwPrismaAirsHandler(CustomGuardrail): @log_guardrail_information async def async_post_call_success_hook( self, - data: dict[str, Any], + data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, - response: Any, - ) -> Any: + response: object, + ) -> object: """ Post-call hook to scan LLM responses before returning to user. @@ -1193,7 +1198,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): assembled_model_response: ModelResponse, request_data: dict, start_time: datetime, - ) -> tuple[bool, ModelResponse, dict[str, Any]]: + ) -> tuple[bool, ModelResponse, dict[str, object]]: """ Scan assembled streaming response and apply masking if needed. Returns (content_was_modified, response, scan_result). @@ -1255,8 +1260,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: Any, - request_data: dict, + response: AsyncIterable[object], + request_data: dict[str, object], ): """ Process streaming response chunks and scan the assembled response. @@ -1367,7 +1372,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # returns a proper JSON error response with the correct status code. # (Raising from a generator hits create_response's generic except → 500.) detail: Final = e.detail if isinstance(e.detail, dict) else {"message": str(e.detail)} - error_obj: Final[dict[str, Any]] = dict(detail.get("error", detail)) + error_obj: Final[dict[str, object]] = dict(detail.get("error", detail)) error_obj["code"] = e.status_code yield f"data: {json.dumps({'error': error_obj})}\n\n" except Exception as e: @@ -1378,8 +1383,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): self, tool_calls: list, is_response: bool, - metadata: dict[str, Any], - call_id: str, + metadata: Mapping[str, object], + call_id: object, request_data: dict, start_time: datetime, ) -> None: @@ -1416,7 +1421,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): tool_name = func.get("name") # --- build tool_event payload (canonical PANW schema) ----------- - tool_event: dict[str, Any] = { + tool_event: dict[str, object] = { "metadata": { "ecosystem": "openai", "method": "tools/call", @@ -1472,7 +1477,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): @staticmethod def _is_anthropic_request( - request_data: dict, + request_data: Mapping[str, object], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> bool: """Detect if the current request is an Anthropic /v1/messages call.""" @@ -1497,7 +1502,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): def _use_latest_user_only( self, - request_data: dict, + request_data: Mapping[str, object], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> bool: """Resolve whether to scan only the latest user message. @@ -1515,8 +1520,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): @staticmethod def _get_latest_user_text_indices( - texts: list[str], - messages: list, + texts: Sequence[str], + messages: Sequence[object], ) -> set | None: """Return text indices belonging to only the latest scannable human-authored (user or developer) message. @@ -1569,8 +1574,8 @@ class PanwPrismaAirsHandler(CustomGuardrail): @staticmethod def _get_scannable_text_indices( - texts: list[str], - structured_messages: list, + texts: Sequence[str], + structured_messages: Sequence[object], ) -> set | None: """Derive which ``texts`` indices originate from user/system messages. @@ -1627,7 +1632,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: @@ -1798,7 +1803,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): # "mcp_tool_name"/"mcp_arguments". Check canonical first, then fallback. mcp_tool_name: Final = request_data.get("mcp_tool_name") or self._mcp_name_fallback(request_data) if mcp_tool_name and input_type == "request": - mcp_tool_event: Final[dict[str, Any]] = { + mcp_tool_event: Final[dict[str, object]] = { "metadata": { "ecosystem": "mcp", "method": "tools/call", diff --git a/litellm/proxy/hooks/mcp_semantic_filter/hook.py b/litellm/proxy/hooks/mcp_semantic_filter/hook.py index 1e57dffa149..3ce406eef73 100644 --- a/litellm/proxy/hooks/mcp_semantic_filter/hook.py +++ b/litellm/proxy/hooks/mcp_semantic_filter/hook.py @@ -5,6 +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 Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Optional from fastapi import HTTPException @@ -164,7 +165,7 @@ class SemanticToolFilterHook(CustomLogger): return [name for name in names if name] @staticmethod - def _narrow_mcp_references(tools: list[Any], selected_tool_names: list[str]) -> list[Any]: + def _narrow_mcp_references(tools: Sequence[Mapping[str, object]], selected_tool_names: list[str]) -> list[object]: """ Restrict each litellm_proxy MCP reference to the semantically selected tools. diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 53c5112d1d7..3fd2adda480 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -31,6 +31,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( + ESTIMATED_OUTPUT_TOKENS_FIELD, + get_estimated_output_tokens, get_key_tag_rpm_limit, get_model_rate_limit_from_metadata, ) @@ -562,6 +564,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data: dict, model: str | None = None, min_configured_tpm_limit: int | None = None, + configured_output_tokens: int | None = None, ) -> int: """ Estimate total tokens this request will consume so we can reserve them @@ -575,6 +578,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): provided, the no-``max_tokens`` output-budget floor is capped at a fraction of that limit so small TPM caps remain usable. Omit to preserve the unconstrained floor. + + ``configured_output_tokens`` is the operator-declared estimate resolved + from key or team metadata. When provided it replaces the heuristic + floor entirely, so the reservation reflects what this tenant's model + actually emits rather than one constant shared by every tenant. """ messages = data.get("messages") prompt: Final = data.get("prompt") @@ -604,7 +612,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): case (_, embeddings_input) if embeddings_input: # Embeddings have no output tokens max_tokens_estimate = 0 - case _ if total_chars == 0: + case _ if total_chars == 0 and configured_output_tokens is None: # Fully contentless request (no messages, prompt, or input). # Don't apply the conservative output-budget floor here — it # would over-reserve and could push small TPM limits into a @@ -619,7 +627,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # so a small per-tenant TPM cap can't be tripped by the floor # alone. output_floor: Final = self._no_max_tokens_output_floor(min_configured_tpm_limit) - max_tokens_estimate = max(estimated_input_tokens, output_floor) + max_tokens_estimate = ( + configured_output_tokens + if configured_output_tokens is not None + else max(estimated_input_tokens, output_floor) + ) total_estimated: Final = estimated_input_tokens + max_tokens_estimate @@ -2586,8 +2598,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data.get("max_tokens") is not None or data.get("max_completion_tokens") is not None ) is_embedding: Final = data.get("input") is not None + configured_output_tokens: Final = get_estimated_output_tokens( + user_api_key_dict=user_api_key_dict, + model_name=requested_model, + ) if capped_floor < baseline_floor and not has_explicit_max_tokens and not is_embedding: - data["max_tokens"] = capped_floor + data["max_tokens"] = max(capped_floor, configured_output_tokens or 0) # Floor at 1 token so contentless requests (/responses, # tool-call continuations, empty messages) still flow @@ -2601,10 +2617,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data=data, model=requested_model, min_configured_tpm_limit=min_configured_tpm_limit, + configured_output_tokens=configured_output_tokens, ), 1, ) + if configured_output_tokens is not None and estimated_tokens > min_configured_tpm_limit: + verbose_proxy_logger.debug( + "Reserving %s tokens for model %s (declared %s=%s plus the input estimate) exceeds the " + "smallest TPM limit this request is charged against (%s), so it cannot be admitted even " + "against an empty window. Lower the declared estimate or raise the TPM limit.", + estimated_tokens, + requested_model, + ESTIMATED_OUTPUT_TOKENS_FIELD, + configured_output_tokens, + min_configured_tpm_limit, + ) + tpm_response: Final = await self.reserve_tpm_tokens( descriptors=descriptors, estimated_tokens=estimated_tokens, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3346f9d7e3b..0e22b5324c1 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -48,6 +48,15 @@ _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( } ) +# Both spellings, because call_type reaches the callback as str(...) of either the +# enum member or its value. +_CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( + ( + CallTypes.aretrieve_batch.value, + str(CallTypes.aretrieve_batch), + ) +) + class _ProxyDBLogger(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -212,7 +221,10 @@ class _ProxyDBLogger(CustomLogger): # Only fetch key details when user_id wasn't already populated (e.g. direct MCP REST calls). # Avoids a cache/DB lookup on every normal LLM request. if metadata.get("user_api_key") and not metadata.get("user_api_key_user_id"): - metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata) + metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( # rebind-ok: enriched metadata replaces the original + metadata=metadata, + resolve_missing_key_identity=str(kwargs.get("call_type")) not in _CAPTURED_IDENTITY_CALL_TYPES, + ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) @@ -337,7 +349,7 @@ class _ProxyDBLogger(CustomLogger): spend_log_error("Error in tracking cost callback - %s", str(e), exc=e) @staticmethod - async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict: + async def _enrich_failure_metadata_with_key_info(metadata: dict, resolve_missing_key_identity: bool = True) -> dict: """ Enriches failure spend log metadata by looking up the key object (and team object) from cache/DB when key fields are missing. @@ -349,6 +361,11 @@ class _ProxyDBLogger(CustomLogger): 2. Post-auth failures (provider errors, rate limits): key fields are populated but team_alias is missing because LiteLLM_VerificationTokenView SQL view doesn't include it. We look up the team object to fill in team_alias. + + Scenario 1 reads the key's identity as it stands right now, so it is only correct + for a log emitted within the request it describes. Callers that log after a delay, + against an identity captured earlier, pass resolve_missing_key_identity=False and + keep their own user_id, team_id and org_id. """ api_key_hash: Final = metadata.get("user_api_key") if not api_key_hash: @@ -361,7 +378,7 @@ class _ProxyDBLogger(CustomLogger): ) # Step 1: If key fields are missing, look up the full key object - if metadata.get("user_api_key_alias") is None: + if resolve_missing_key_identity and metadata.get("user_api_key_alias") is None: try: key_obj: Final = await get_key_object( hashed_token=api_key_hash, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index c48fee96646..10142a894a1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -5,6 +5,7 @@ import re import time from collections import OrderedDict from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final from fastapi import HTTPException, Request @@ -66,6 +67,32 @@ _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") _SHA256_HEX_RE: Final = re.compile(r"^[0-9a-f]{64}$") +# W3C Trace Context traceparent header: https://www.w3.org/TR/trace-context/ +# e.g. "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" +_TRACEPARENT_RE: Final = re.compile(r"^[0-9a-f]{2}-([0-9a-f]{32})-[0-9a-f]{16}-[0-9a-f]{2}$", re.IGNORECASE) + + +def _trace_id_from_traceparent(traceparent: str) -> str | None: + """Extract the trace-id from a W3C Trace Context traceparent header, e.g. + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" -> the 32-hex + trace-id in the middle. An all-zero trace-id is invalid per spec and is + rejected, matching how the OpenTelemetry SDK itself treats it.""" + match: Final = _TRACEPARENT_RE.match(traceparent.strip()) + if not match: + return None + trace_id: Final = match.group(1).lower() + return trace_id if trace_id != "0" * 32 else None + + +def _session_id_from_baggage(baggage: str) -> str | None: + """Extract a session.id entry from a W3C Baggage header + (https://www.w3.org/TR/baggage/), e.g. "session.id=abc-123,user.id=42".""" + for pair in baggage.split(","): + key, _, value = pair.strip().partition("=") + if key.strip() == "session.id" and value.strip(): + return value.strip() + return None + def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: """Only proxy-validated keys are stamped, proven by the unforgeable @@ -210,6 +237,11 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_session_scoped", "max_agentic_loops", + # Recomputed below from the actual caller-controlled timeout sources (headers and + # body fields); a client-forged value here would let a request either dodge cooldown + # protection on a real deployment failure or force a false "not caller-controlled" + # reading that lets its own bad timeout cool down deployments other tenants rely on. + "client_side_timeout", ) _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = ( @@ -850,6 +882,19 @@ class LiteLLMProxyRequestSetup: return float(stream_timeout_header) return None + @staticmethod + def _get_keepalive_seconds_from_request(headers: Mapping[str, str]) -> float | None: + """ + Get `keepalive_seconds` from the request headers, for clients (e.g. the + Vercel AI SDK) that can set custom headers more easily than extra body + fields. Subject to the same deployment-level allow_client_keepalive_override + gate as the request body field: see _resolve_keepalive_seconds. + """ + keepalive_seconds_header: Final = headers.get("x-litellm-keepalive-seconds", None) + if keepalive_seconds_header is not None: + return float(keepalive_seconds_header) + return None + @staticmethod def _get_num_retries_from_request(headers: dict) -> int | None: """ @@ -1035,6 +1080,7 @@ class LiteLLMProxyRequestSetup: def add_litellm_data_for_backend_llm_call( *, headers: dict, + request_data: Mapping[str, Any], user_api_key_dict: UserAPIKeyAuth, general_settings: dict[str, Any] | None = None, ) -> LitellmDataForBackendLLMCall: @@ -1053,18 +1099,38 @@ class LiteLLMProxyRequestSetup: if _organization is not None: data["organization"] = _organization - timeout: Final = LiteLLMProxyRequestSetup._get_timeout_from_request(headers) - if timeout is not None: - data["timeout"] = timeout + header_timeout: Final = LiteLLMProxyRequestSetup._get_timeout_from_request(headers) + if header_timeout is not None: + data["timeout"] = header_timeout - stream_timeout: Final = LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers) - if stream_timeout is not None: - data["stream_timeout"] = stream_timeout + header_stream_timeout: Final = LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers) + if header_stream_timeout is not None: + data["stream_timeout"] = header_stream_timeout + + # Router._get_timeout resolves the effective per-attempt timeout from any of + # kwargs["timeout"], kwargs["request_timeout"], or kwargs["stream_timeout"], and a + # caller can supply any of those via the request body as well as the headers above. + # A deliberately tiny value can force a 408 on every deployment in a fallback chain, + # so this marker (never trusted verbatim from the client; stripped above) must cover + # every source cooldown_handlers._trigger_cooldown_for_failed_deployment needs to + # distinguish from a real deployment health signal. + if ( + header_timeout is not None + or header_stream_timeout is not None + or request_data.get("timeout") is not None + or request_data.get("request_timeout") is not None + or request_data.get("stream_timeout") is not None + ): + data["client_side_timeout"] = True num_retries: Final = LiteLLMProxyRequestSetup._get_num_retries_from_request(headers) if num_retries is not None: data["num_retries"] = num_retries + keepalive_seconds: Final = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request(headers) + if keepalive_seconds is not None: + data["keepalive_seconds"] = keepalive_seconds + return data @staticmethod @@ -1113,6 +1179,33 @@ class LiteLLMProxyRequestSetup: body_metadata["user_id"] = session_id verbose_proxy_logger.debug("Extracted session_id from Anthropic metadata.user_id") + # Last-resort fallback: the W3C standards for trace/session propagation + # (https://www.w3.org/TR/trace-context/, https://www.w3.org/TR/baggage/). + # Lower priority than everything above - only fires when neither the + # explicit litellm headers nor the Anthropic-metadata path found + # anything - but lets a caller's existing traceparent/baggage headers + # (from real OTel instrumentation) correlate with litellm's own logs + # instead of generating an unrelated trace_id. + normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)}) + if "litellm_trace_id" not in data: + traceparent: Final = normalized_headers.get("traceparent") + if isinstance(traceparent, str): + trace_id_from_traceparent: Final = _trace_id_from_traceparent(traceparent) + if trace_id_from_traceparent: + metadata_from_headers["trace_id"] = trace_id_from_traceparent + data["litellm_trace_id"] = trace_id_from_traceparent # rebind-ok: data is an out-param + verbose_proxy_logger.debug( + "Extracted trace_id from W3C traceparent header: %s", trace_id_from_traceparent + ) + if "litellm_session_id" not in data: + baggage: Final = normalized_headers.get("baggage") + if isinstance(baggage, str): + session_id_from_baggage: Final = _session_id_from_baggage(baggage) + if session_id_from_baggage: + metadata_from_headers["session_id"] = session_id_from_baggage + data["litellm_session_id"] = session_id_from_baggage # rebind-ok: data is an out-param + verbose_proxy_logger.debug("Extracted session_id from W3C baggage header") + if isinstance(data[_metadata_variable_name], dict): data[_metadata_variable_name].update(metadata_from_headers) return data @@ -1545,6 +1638,7 @@ async def add_litellm_data_to_request( data.update( LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( headers=_headers, + request_data=data, user_api_key_dict=user_api_key_dict, general_settings=general_settings, ) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 446ea76752e..8c6195388c5 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -20,7 +20,10 @@ from fastapi import APIRouter, Depends, HTTPException from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.proxy.management_endpoints.common_utils import ( + _user_has_admin_view, + validate_budget_duration, +) from litellm.proxy.utils import jsonify_object from litellm.repositories.budget_repository import BudgetRepository @@ -72,6 +75,8 @@ async def new_budget( detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"}, ) + validate_budget_duration(budget_obj.budget_duration) + # Validate model_max_budget if present if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0: from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -153,6 +158,8 @@ async def update_budget( detail={"error": f"soft_budget must be a non-negative finite number. Received: {budget_obj.soft_budget}"}, ) + validate_budget_duration(budget_obj.budget_duration) + # Validate model_max_budget if present in update if budget_obj.model_max_budget is not None and len(budget_obj.model_max_budget) > 0: from litellm.proxy.management_endpoints.key_management_endpoints import ( diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index 50637208e03..53d03bc7ba6 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -12,7 +12,7 @@ import asyncio import json from collections.abc import Mapping from datetime import datetime, timezone -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import APIRouter, Depends, Header, HTTPException from pydantic import BaseModel, Field @@ -37,8 +37,26 @@ from litellm.types.management_endpoints import ( CacheSettingsField, ) +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + router: Final = APIRouter() + +class _CacheConfigRow(Protocol): + cache_settings: str | Mapping[str, object] | None + + +class _CacheConfigTable(Protocol): + async def find_unique(self, where: Mapping[str, str]) -> _CacheConfigRow | None: ... + + async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> _CacheConfigRow: ... + + +def _cache_config_table(prisma_client: "PrismaClient") -> _CacheConfigTable: + return CacheConfigRepository(prisma_client).table + + # Cache fields holding credentials. Masked on read so plaintext Redis / # Sentinel passwords never leave the server in a GET response. `url` is here # because a Redis/Valkey URL can embed a password inline @@ -197,7 +215,7 @@ def _saved_secret_is_reusable(incoming: Mapping[str, object], saved: Mapping[str return True -def _merge_over_saved(incoming: Mapping[str, object], saved: Mapping[str, object]) -> dict[str, Any]: +def _merge_over_saved(incoming: Mapping[str, object], saved: Mapping[str, object]) -> Mapping[str, object]: """Keep the stored secret behind any credential the caller echoed back redacted or omitted. GET returns credentials as the marker and the form never re-prefills a @@ -339,7 +357,7 @@ class CacheSettingsManager: return normalized1 == normalized2 @staticmethod - async def init_cache_settings_in_db(prisma_client, proxy_config): + async def init_cache_settings_in_db(prisma_client: "PrismaClient", proxy_config): """ Initialize cache settings from database into the router on startup. Only reinitializes if cache params have changed. @@ -349,7 +367,7 @@ class CacheSettingsManager: try: cache_config: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}), + lambda: _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"}), reason="init_cache_settings_in_db_lookup_failure", ) if cache_config is not None and cache_config.cache_settings: @@ -444,7 +462,7 @@ async def get_cache_settings( # Read the stored settings (decrypted); an env-only cache has none. stored: dict[str, object] = {} if prisma_client is not None: - cache_config = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}) + cache_config = await _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"}) if cache_config is not None and cache_config.cache_settings: stored = proxy_config._decrypt_db_variables( variables_dict=_parse_stored_settings(cache_config.cache_settings) @@ -511,9 +529,7 @@ async def test_cache_connection( saved_settings: dict[str, object] = {} if prisma_client is not None: try: - existing_row: Final = await CacheConfigRepository(prisma_client).table.find_unique( - where={"id": "cache_config"} - ) + existing_row: Final = await _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"}) if existing_row is not None and existing_row.cache_settings: saved_settings = proxy_config._decrypt_db_variables( variables_dict=_parse_stored_settings(existing_row.cache_settings) @@ -590,7 +606,7 @@ async def update_cache_settings( try: # Read the stored row first: its decrypted values back any credential the # caller echoed back redacted, and its key set drives the audit diff. - existing_row: Final = await CacheConfigRepository(prisma_client).table.find_unique(where={"id": "cache_config"}) + existing_row: Final = await _cache_config_table(prisma_client).find_unique(where={"id": "cache_config"}) before_settings: dict[str, object] | None = None saved_settings: dict[str, object] = {} if existing_row is not None and existing_row.cache_settings: @@ -606,7 +622,7 @@ async def update_cache_settings( encrypted_settings: Final = proxy_config._encrypt_env_variables(environment_variables=cache_settings) # Save to database - await CacheConfigRepository(prisma_client).table.upsert( + await _cache_config_table(prisma_client).upsert( where={"id": "cache_config"}, data={ "create": { diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 7a30f6b799a..d1542b38996 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -8,7 +8,9 @@ from fastapi import HTTPException, status from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger +from litellm.constants import PTU_SENTINEL_API_KEY from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import DeletedVerificationTokenRepository from litellm.repositories.verification_token_repository import ( @@ -140,6 +142,28 @@ class _GroupingSetsRow(SimpleNamespace): failed_requests: int | None +def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float: + """Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled. + + Both read paths funnel through here: the paginated path reads the ``ptu_flat_cost`` + column straight off the row, and the aggregated path reads the SUM() alias. Rows an + operator accrued during an earlier opt-in stay in the table, so the gate lives on the + read rather than on the query that produced the rows. + + The row is checked before the flag because this runs once per metric accumulation, and + a record fans out across roughly a dozen breakdowns. The flag reads through the secret + manager, uncached, so consulting it for every accumulation put thousands of lookups on + a shared endpoint that made none before. Only a row actually carrying flat cost, which + is a sentinel row, reaches it now. + """ + raw: Final = getattr(record, "ptu_flat_cost", None) or 0.0 + if not raw: + return 0.0 + if not is_ptu_cost_attribution_enabled(): + return 0.0 + return raw + + def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) -> SpendMetrics: """Update metrics with new record data. @@ -150,6 +174,7 @@ def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) -> prompt_tokens: Final = record.prompt_tokens or 0 completion_tokens: Final = record.completion_tokens or 0 existing_metrics.spend += record.spend or 0.0 + existing_metrics.flat_cost += _reported_flat_cost(record) existing_metrics.prompt_tokens += prompt_tokens existing_metrics.completion_tokens += completion_tokens existing_metrics.total_tokens += prompt_tokens + completion_tokens @@ -208,30 +233,43 @@ def update_breakdown_metrics( entity_id_field: str | None = None, entity_metadata_field: Mapping[str, dict[str, object]] | None = None, ) -> BreakdownMetrics: - """Updates breakdown metrics for a single record using the existing update_metrics function""" + """Updates breakdown metrics for a single record using the existing update_metrics function. + + PTU sentinel rows (api_key == PTU_SENTINEL_API_KEY) add their flat cost to every + parent bucket but never appear as an api_key row, and are kept out of the + per-request provider breakdown.""" + + is_ptu_sentinel: Final = record.api_key == PTU_SENTINEL_API_KEY + + # A PTU sentinel row keys on the deployment id so a rename cannot move it, and carries + # the operator-facing name in model_group. The breakdown key is rendered directly as a + # label, so display the name; two deployments sharing one name merge here, which is + # what the write path used to do by collapsing them into a single row. + model_key: Final = (record.model_group or record.model) if is_ptu_sentinel else record.model # Update model breakdown - if record.model and record.model not in breakdown.models: - breakdown.models[record.model] = MetricWithMetadata( + if model_key and model_key not in breakdown.models: + breakdown.models[model_key] = MetricWithMetadata( metrics=SpendMetrics(), - metadata=model_metadata.get(record.model, {}), # Add any model-specific metadata here + metadata=model_metadata.get(model_key, {}), # Add any model-specific metadata here ) - if record.model: - breakdown.models[record.model].metrics = update_metrics(breakdown.models[record.model].metrics, record) + if model_key: + breakdown.models[model_key].metrics = update_metrics(breakdown.models[model_key].metrics, record) - # Update API key breakdown for this model - if record.api_key not in breakdown.models[record.model].api_key_breakdown: - breakdown.models[record.model].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + if not is_ptu_sentinel: + # Update API key breakdown for this model + if record.api_key not in breakdown.models[model_key].api_key_breakdown: + breakdown.models[model_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), + ) + breakdown.models[model_key].api_key_breakdown[record.api_key].metrics = update_metrics( + breakdown.models[model_key].api_key_breakdown[record.api_key].metrics, + record, ) - breakdown.models[record.model].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.models[record.model].api_key_breakdown[record.api_key].metrics, - record, - ) # Update model group breakdown model_group_key: Final = record.model_group or record.model @@ -245,19 +283,20 @@ def update_breakdown_metrics( breakdown.model_groups[model_group_key].metrics, record ) - # Update API key breakdown for this model - if record.api_key not in breakdown.model_groups[model_group_key].api_key_breakdown: - breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + if not is_ptu_sentinel: + # Update API key breakdown for this model + if record.api_key not in breakdown.model_groups[model_group_key].api_key_breakdown: + breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), + ) + breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics = update_metrics( + breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics, + record, ) - breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics, - record, - ) if record.mcp_namespaced_tool_name: if record.mcp_namespaced_tool_name not in breakdown.mcp_servers: @@ -288,28 +327,29 @@ def update_breakdown_metrics( record, ) - # Update provider breakdown - provider: Final = record.custom_llm_provider or "unknown" - if provider not in breakdown.providers: - breakdown.providers[provider] = MetricWithMetadata( - metrics=SpendMetrics(), - metadata=provider_metadata.get(provider, {}), # Add any provider-specific metadata here - ) - breakdown.providers[provider].metrics = update_metrics(breakdown.providers[provider].metrics, record) + if not is_ptu_sentinel: + # Update provider breakdown + provider: Final = record.custom_llm_provider or "unknown" + if provider not in breakdown.providers: + breakdown.providers[provider] = MetricWithMetadata( + metrics=SpendMetrics(), + metadata=provider_metadata.get(provider, {}), # Add any provider-specific metadata here + ) + breakdown.providers[provider].metrics = update_metrics(breakdown.providers[provider].metrics, record) - # Update API key breakdown for this provider - if record.api_key not in breakdown.providers[provider].api_key_breakdown: - breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + # Update API key breakdown for this provider + if record.api_key not in breakdown.providers[provider].api_key_breakdown: + breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), + ) + breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics( + breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, + record, ) - breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.providers[provider].api_key_breakdown[record.api_key].metrics, - record, - ) # Update endpoint breakdown if record.endpoint: @@ -336,16 +376,17 @@ def update_breakdown_metrics( record, ) - # Update api key breakdown - if record.api_key not in breakdown.api_keys: - breakdown.api_keys[record.api_key] = KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), # Add any api_key-specific metadata here - ) - breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record) + if not is_ptu_sentinel: + # Update api key breakdown + if record.api_key not in breakdown.api_keys: + breakdown.api_keys[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), # Add any api_key-specific metadata here + ) + breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record) # Update entity-specific metrics if entity_id_field is provided if entity_id_field: @@ -358,19 +399,20 @@ def update_breakdown_metrics( ) breakdown.entities[entity_value].metrics = update_metrics(breakdown.entities[entity_value].metrics, record) - # Update API key breakdown for this entity - if record.api_key not in breakdown.entities[entity_value].api_key_breakdown: - breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( - metrics=SpendMetrics(), - metadata=KeyMetadata( - key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), - team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), - ), + if not is_ptu_sentinel: + # Update API key breakdown for this entity + if record.api_key not in breakdown.entities[entity_value].api_key_breakdown: + breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata( + metrics=SpendMetrics(), + metadata=KeyMetadata( + key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None), + team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None), + ), + ) + breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics( + breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics, + record, ) - breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics( - breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics, - record, - ) return breakdown @@ -599,6 +641,14 @@ def _build_aggregated_sql_query( # total_successful_requests metadata they feed) once the admin UI reads SGR # only from LiteLLM_DailyGatewayRequests. The remaining spend, token and # api_requests rollups are still served from here. + # + # Only LiteLLM_DailyTeamSpend carries ptu_flat_cost; other daily tables emit a + # constant zero so the SpendMetrics.flat_cost response shape stays uniform. + ptu_flat_cost_select: Final = ( + "SUM(ptu_flat_cost)::float AS ptu_flat_cost" + if table_name == "litellm_dailyteamspend" + else "0::float AS ptu_flat_cost" + ) sql_query: Final = f""" SELECT date, @@ -612,6 +662,7 @@ def _build_aggregated_sql_query( custom_llm_provider, mcp_namespaced_tool_name, endpoint) AS group_level, SUM(spend)::float AS spend, + {ptu_flat_cost_select}, SUM(prompt_tokens)::bigint AS prompt_tokens, SUM(completion_tokens)::bigint AS completion_tokens, SUM(cache_read_input_tokens)::bigint AS cache_read_input_tokens, @@ -707,7 +758,9 @@ async def _aggregate_spend_records( The per-row loop is offloaded to a worker thread via asyncio.to_thread so a large result set doesn't peg the event loop. """ - api_keys: Final[set[str]] = {record.api_key for record in records if record.api_key} + api_keys: Final[set[str]] = { + record.api_key for record in records if record.api_key and record.api_key != PTU_SENTINEL_API_KEY + } api_key_metadata: dict[str, _KeyMetadataDict] = {} if api_keys: @@ -754,6 +807,7 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics: completion_tokens: Final = record.completion_tokens or 0 return SpendMetrics( spend=record.spend or 0.0, + flat_cost=_reported_flat_cost(record), prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, @@ -820,6 +874,7 @@ def _aggregate_grouping_sets_records_sync( for record in records: level = record.group_level metrics = _record_to_spend_metrics(record) + is_ptu_sentinel = record.api_key == PTU_SENTINEL_API_KEY if level == _GROUP_GRAND_TOTAL: total_metrics = metrics @@ -832,7 +887,7 @@ def _aggregate_grouping_sets_records_sync( breakdown = ensure_date(record.date)["breakdown"] if level == _GROUP_DATE_API_KEY: - if record.api_key: + if record.api_key and not is_ptu_sentinel: breakdown.api_keys[record.api_key] = KeyMetricWithMetadata( metrics=metrics, metadata=_key_metadata(api_key_metadata, record.api_key), @@ -841,13 +896,13 @@ def _aggregate_grouping_sets_records_sync( if record.model: assign_metric_with_metadata(breakdown.models, record.model, metrics) elif level == _GROUP_DATE_MODEL_API_KEY: - if record.model and record.api_key: + if record.model and record.api_key and not is_ptu_sentinel: assign_api_key_breakdown(breakdown.models, record.model, record.api_key, metrics) elif level == _GROUP_DATE_MODEL_GROUP: if record.model_group: assign_metric_with_metadata(breakdown.model_groups, record.model_group, metrics) elif level == _GROUP_DATE_MODEL_GROUP_API_KEY: - if record.model_group and record.api_key: + if record.model_group and record.api_key and not is_ptu_sentinel: assign_api_key_breakdown( breakdown.model_groups, record.model_group, @@ -855,10 +910,17 @@ def _aggregate_grouping_sets_records_sync( metrics, ) elif level == _GROUP_DATE_PROVIDER: + # Only PTU sentinel rows carry ptu_flat_cost and they have no provider, so at + # this level the sentinel's cost would land under "unknown". Withholding the + # flat cost matches the per-row path, which skips sentinel rows outright. The + # bucket itself is still assigned unconditionally: a legacy row predating the + # api_requests column backfills to all zeroes, and skipping those would drop a + # provider the base build reported. + provider_metrics = metrics.model_copy(update={"flat_cost": 0.0}) # mutable-ok: pydantic update payload provider = record.custom_llm_provider or "unknown" - assign_metric_with_metadata(breakdown.providers, provider, metrics) + assign_metric_with_metadata(breakdown.providers, provider, provider_metrics) elif level == _GROUP_DATE_PROVIDER_API_KEY: - if record.api_key: + if record.api_key and not is_ptu_sentinel: provider = record.custom_llm_provider or "unknown" assign_api_key_breakdown(breakdown.providers, provider, record.api_key, metrics) elif level == _GROUP_DATE_MCP: @@ -898,7 +960,7 @@ async def _aggregate_grouping_sets_records( records: Sequence[_GroupingSetsRow], ) -> _AggregatedSpendData: """Async wrapper: fetch api_key_metadata, then dispatch on a worker thread.""" - api_keys: Final[set[str]] = {r.api_key for r in records if r.api_key} + api_keys: Final[set[str]] = {r.api_key for r in records if r.api_key and r.api_key != PTU_SENTINEL_API_KEY} api_key_metadata: dict[str, _KeyMetadataDict] = {} if api_keys: @@ -1008,6 +1070,7 @@ async def get_daily_activity( results=aggregated["results"], metadata=DailySpendMetadata( total_spend=metadata_metrics.spend, + total_flat_cost=metadata_metrics.flat_cost, total_prompt_tokens=metadata_metrics.prompt_tokens, total_completion_tokens=metadata_metrics.completion_tokens, total_tokens=metadata_metrics.total_tokens, @@ -1098,6 +1161,7 @@ async def get_daily_activity_aggregated( results=aggregated["results"], metadata=DailySpendMetadata( total_spend=aggregated["totals"].spend, + total_flat_cost=aggregated["totals"].flat_cost, total_prompt_tokens=aggregated["totals"].prompt_tokens, total_completion_tokens=aggregated["totals"].completion_tokens, total_tokens=aggregated["totals"].total_tokens, diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 3868b04f385..2241884faf1 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -22,6 +22,35 @@ def validate_finite_spend(spend: float | None) -> None: ) +def validate_budget_duration(budget_duration: str | None) -> None: + """Reject budget durations that can't be parsed, are non-positive, or + overflow date math, so a bad value can't be persisted and later crash the + budget reset job. + + A non-positive duration also resolves to a reset time of "now", which leaves + the row permanently due: the reset job re-reads it every tick and, once + enough of them exist, they fill each batch and starve every other tenant's + reset. + """ + if budget_duration is None: + return + + from litellm.litellm_core_utils.duration_parser import duration_in_seconds + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + try: + if duration_in_seconds(budget_duration) <= 0: + raise ValueError("budget_duration must be positive") + get_budget_reset_time(budget_duration=budget_duration) + except (ValueError, OverflowError): + raise HTTPException( + status_code=400, + detail={ + "error": f"Invalid budget_duration '{budget_duration}'. Use a format like '1h', '24h', '7d', or '30d'." + }, + ) + + from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index a51ff48aab6..bfc70da46ea 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.management_endpoints.common_utils import validate_budget_duration from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, handle_update_object_permission_common, @@ -184,6 +185,7 @@ def new_budget_request(data: NewCustomerRequest) -> BudgetNewRequest | None: if budget_kv_pairs: budget_request: Final = BudgetNewRequest(**budget_kv_pairs) + validate_budget_duration(budget_request.budget_duration) if budget_request.budget_reset_at is None and budget_request.budget_duration is not None: budget_request.budget_reset_at = datetime.utcnow() + timedelta( seconds=duration_in_seconds(duration=budget_request.budget_duration) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index abc5d3e53ff..a416a197ab8 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -42,6 +42,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _user_has_admin_view, require_caller_user_id_for_non_admin, + validate_budget_duration, validate_finite_spend, ) from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -506,6 +507,8 @@ async def new_user( status_code=500, detail=CommonProxyErrors.db_not_connected_error.value, ) + validate_budget_duration(data.budget_duration) + # Check for duplicate user_id or email await _check_duplicate_user_id(data.user_id, prisma_client) await _check_duplicate_user_email(data.user_email, prisma_client) @@ -1185,6 +1188,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda if "budget_duration" in non_default_values: from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + validate_budget_duration(non_default_values["budget_duration"]) non_default_values["budget_reset_at"] = get_budget_reset_time( budget_duration=non_default_values["budget_duration"] ) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2de1d177b33..38b5d755535 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -55,7 +55,10 @@ from litellm.proxy.auth.auth_checks import ( get_project_object, get_team_object, ) -from litellm.proxy.auth.auth_utils import abbreviate_api_key +from litellm.proxy.auth.auth_utils import ( + abbreviate_api_key, + enforce_output_token_estimates_are_admin_only, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, @@ -79,6 +82,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _set_object_metadata_field, _team_member_has_permission, _user_has_admin_view, + validate_budget_duration, validate_finite_spend, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( @@ -841,12 +845,21 @@ async def _common_key_generation_helper( premium_user=premium_user, ) + validate_budget_duration(data.budget_duration) + if data.throttle_on_budget_exceeded is True and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."}, ) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) + if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None: await validate_team_id_used_in_service_account_request( team_id=data.team_id, @@ -1014,7 +1027,7 @@ async def _common_key_generation_helper( # Only set budget_duration on key when explicitly provided. Keys with budget_id # but no explicit budget_duration follow their linked budget tier's schedule; - # reset_budget_for_keys_linked_to_budgets() resets them when the tier resets. + # reset_budget_for_litellm_budget_table() resets them when the tier resets. # This avoids duplicating budget_duration on keys so tier updates apply automatically. if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) @@ -1584,6 +1597,8 @@ async def generate_key_fn( - 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"]}. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. + - default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer. Falls back to the team setting, then to the built-in estimate. + - default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above. Example - {"gpt-4": 4096, "gpt-3.5-turbo": 1024}. Takes precedence over the key-wide value. - mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit. - tag_rpm_limit: Optional[dict] - key-specific per-request-tag rpm limit, keyed by request tag. Example - {"cell-1": 1000, "cell-2": 500}. Each tag gets an independent counter; requests whose tag is absent fall back to the key-level rpm limit. - tpm_limit_type: Optional[str] - Type of tpm limit. Options: "best_effort_throughput" (no error if we're overallocating tpm), "guaranteed_throughput" (raise an error if we're overallocating tpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput". @@ -1793,6 +1808,8 @@ async def generate_service_account_key_fn( - 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"]}. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. + - default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer. Falls back to the team setting, then to the built-in estimate. + - default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above. Example - {"gpt-4": 4096, "gpt-3.5-turbo": 1024}. Takes precedence over the key-wide value. - mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit. - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" @@ -2387,6 +2404,7 @@ async def _validate_update_key_data( """Validate permissions and constraints for key update.""" # Reject NaN/±inf spend before it can reach the DB / spend counter. validate_finite_spend(data.spend) + validate_budget_duration(data.budget_duration) _is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value @@ -2473,6 +2491,13 @@ async def _validate_update_key_data( detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."}, ) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) + # Personal-key bypass: the caller both created the key AND still owns it # (user_id == caller). Checking only created_by would let a demoted admin # who originally created a key for another user continue editing it without @@ -2655,6 +2680,8 @@ async def update_key_fn( - mcp_rpm_limit: Optional[dict] - Per-MCP-server RPM limits, keyed by MCP server name {"github": 100, "slack": 200} - tag_rpm_limit: Optional[dict] - Per-request-tag RPM limits, keyed by request tag {"cell-1": 1000, "cell-2": 500}. Each tag gets an independent counter; absent tags fall back to the key-level rpm limit. - model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000} + - default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer. + - default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above {"gpt-4": 4096, "gpt-3.5-turbo": 1024} - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - allowed_cache_controls: Optional[list] - List of allowed cache control values @@ -4629,6 +4656,15 @@ async def _execute_virtual_key_regeneration( prisma_client=prisma_client, ) + if data is not None: + _existing_key_metadata: Final = getattr(key_in_db, "metadata", None) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) + new_token: Final = await get_new_token(data=data) new_token_hash: Final = hash_token(new_token) new_token_key_name: Final = abbreviate_api_key(api_key=new_token) diff --git a/litellm/proxy/management_endpoints/management_v1/common.py b/litellm/proxy/management_endpoints/management_v1/common.py index 8525d67a041..ec79820465a 100644 --- a/litellm/proxy/management_endpoints/management_v1/common.py +++ b/litellm/proxy/management_endpoints/management_v1/common.py @@ -4,7 +4,8 @@ from typing import Final from urllib.parse import urlencode from fastapi import Request -from fastapi.dependencies.utils import get_flat_dependant +from fastapi.dependencies.utils import get_flat_params +from fastapi.params import ParamTypes from fastapi.responses import JSONResponse from litellm.types.proxy.management_endpoints.management_v1 import ( @@ -42,7 +43,13 @@ def _declared_query_params(request: Request) -> frozenset[str]: dependant: Final = getattr(route, "dependant", None) if dependant is None: return frozenset() - return frozenset(field.alias for field in get_flat_dependant(dependant, skip_repeats=True).query_params) + # fastapi>=0.140.7 removed get_flat_dependant(); get_flat_params() returns the + # flattened (deduped) param list. Filter to query params to match the old behavior. + return frozenset( + field.alias + for field in get_flat_params(dependant) + if getattr(field.field_info, "in_", None) == ParamTypes.query + ) def escape_like(value: str) -> str: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index c2087005863..8a52b0d1abb 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -15,6 +15,7 @@ import datetime import json from collections.abc import Awaitable, Mapping, Sequence from json import JSONDecodeError +from types import MappingProxyType from typing import Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -55,6 +56,10 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team as _legacy_update_team, ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log +from litellm.proxy.spend_tracking.ptu_feature_flag import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + is_ptu_cost_attribution_enabled, +) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ModelTableRepository @@ -79,6 +84,7 @@ from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, GenericLiteLLMParams, + ModelInfo, updateDeployment, ) from litellm.utils import get_utc_datetime @@ -233,7 +239,129 @@ def _raise_on_strategy_router_write_violation( ) +_PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to") + + +def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[str]: + """The PTU fields a patch sends as an explicit null, which update_db_model drops. + + Empty while the feature is off, so disabling pauses PTU rather than letting a client + that round-trips a model_info blob erase a configuration set up during an earlier opt-in. + """ + if model_info is None or not is_ptu_cost_attribution_enabled(): + return frozenset() + return frozenset( + field + for field in _PTU_MODEL_INFO_FIELDS + if field in model_info.model_fields_set and getattr(model_info, field) is None + ) + + +def _merged_ptu_model_info(*, db_model: Deployment, patch_data: updateDeployment) -> Mapping[str, object]: + """The model_info a patch would store, which is the stored blob updated by the patch. + + A PTU invariant holds over the deployment as it will exist, not over whichever subset + of fields a caller happened to send. + """ + empty: Final[Mapping[str, object]] = MappingProxyType({}) + stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else empty + incoming: Final = patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else empty + cleared: Final = _explicitly_cleared_ptu_fields(patch_data.model_info) + return MappingProxyType({k: v for k, v in {**stored, **incoming}.items() if k not in cleared}) + + +def _raise_if_ptu_cost_attribution_disabled(incoming_model_info: Mapping[str, object]) -> None: + """Reject PTU model_info fields unless the operator opted into PTU cost attribution. + + Takes the incoming request's model_info rather than the merged deployment, so an + unrelated patch of a model that still stores PTU config from an earlier opt-in is + left alone. The fields are rejected rather than dropped so a caller never believes + a flat cost was configured while the rollup that would price it is not running. + + Only a value is rejected. An explicit null reaches the clear loop, which is gated on + the same flag, so a disabled proxy neither writes PTU config nor erases what an + earlier opt-in stored. Disabling pauses the feature rather than discarding its setup. + """ + if is_ptu_cost_attribution_enabled(): + return + supplied: Final = tuple(field for field in _PTU_MODEL_INFO_FIELDS if incoming_model_info.get(field) is not None) + if not supplied: + return + raise HTTPException( + status_code=400, + detail=( + f"PTU cost attribution is disabled, so {', '.join(supplied)} cannot be set. " + f"Set {PTU_COST_ATTRIBUTION_ENV_VAR}=true to enable it." + ), + ) + + +def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: + """Enforce the PTU cross-field invariant on the effective model_info. + + ptu_count and cost_per_ptu_per_hour must be set together, and a team_id and a + ptu_effective_from are required when they are. The start is mandatory rather than + defaulted because flat cost accrues from it: inferring one would let a deployment + configured today be billed for days it did not exist. Per-field bounds (positive + count, non-negative rate) are enforced by ModelInfo itself. + + Window ordering is checked before the count/rate gate. A patch that touches only one + end of the window carries no count or rate, and ModelInfo sees one field at a time, so + leaving it to either would let an inverted window reach the row; the next load then + fails to parse it and drops the deployment out of the router, where no further patch + can repair it because each one re-parses the stored value first. + """ + effective_from: Final = _coerce_ptu_datetime(model_info.get("ptu_effective_from")) + effective_to: Final = _coerce_ptu_datetime(model_info.get("ptu_effective_to")) + if effective_from is not None and effective_to is not None and effective_to <= effective_from: + raise HTTPException(status_code=400, detail="ptu_effective_to must be after ptu_effective_from") + + has_count: Final = model_info.get("ptu_count") is not None + has_rate: Final = model_info.get("cost_per_ptu_per_hour") is not None + if not has_count and not has_rate: + return + if has_count != has_rate: + raise HTTPException(status_code=400, detail="ptu_count and cost_per_ptu_per_hour must be set together") + if effective_from is None: + raise HTTPException( + status_code=400, + detail=( + "ptu_effective_from is required when PTU fields are set. Flat cost accrues from that " + "instant, so without it the start would have to be inferred and a deployment configured " + "today could be billed for days it did not exist" + ), + ) + if not model_info.get("team_id"): + raise HTTPException( + status_code=400, detail="team_id is required when PTU fields are set (one model maps to one team)" + ) + + +def _parse_ptu_datetime(value: object) -> datetime.datetime | None: + """``value`` as a datetime, parsing an ISO string, else None.""" + if isinstance(value, datetime.datetime): + return value + if not isinstance(value, str): + return None + try: + return datetime.datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + + +def _coerce_ptu_datetime(value: object) -> datetime.datetime | None: + """Coerce a model_info effective-window value (datetime or ISO string) to UTC, else None.""" + parsed: Final = _parse_ptu_datetime(value) + if parsed is None: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=datetime.timezone.utc) + return parsed.astimezone(datetime.timezone.utc) + + def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: + if updated_patch.model_info is not None: + _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) merged_model_name: Final = updated_patch.model_name or db_model.model_name merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True) merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True) @@ -270,6 +398,10 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: merged_model_info.pop(field, None) merged_litellm_params.pop(field, None) + for field in _explicitly_cleared_ptu_fields(updated_patch.model_info): + merged_model_info.pop(field, None) + + _validate_ptu_model_info(merged_model_info) # convert to prisma compatible format @@ -716,6 +848,18 @@ async def _update_team_model_in_db( premium_user=premium_user, ) + # Validated before any write, beside the premium check the create path already runs + # here. The team ACL is updated below and autocommits, so a validator that raises + # further down would leave the team mutated and the deployment row never written. + # + # The merged view is what gets stored, so that is what has to satisfy the invariants. + # Validating the patch alone rejected a partial edit of an already valid deployment: + # raising the rate on a configured model carries no ptu_effective_from, which the + # stored row supplies. + if patch_data.model_info is not None: + _raise_if_ptu_cost_attribution_disabled(patch_data.model_info.model_dump(exclude_none=True)) + _validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data)) + patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None # No team_id in patch, proceed with standard update @@ -1424,6 +1568,10 @@ async def add_new_model( model_response: LiteLLM_ProxyModelTable | None = None # update DB + incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True) + _raise_if_ptu_cost_attribution_disabled(incoming_model_info) + _validate_ptu_model_info(incoming_model_info) + if store_model_in_db is True: """ - store model_list in db diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 97f494c51de..60d3d650d00 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -82,6 +82,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_object, get_user_object, ) +from litellm.proxy.auth.auth_utils import enforce_output_token_estimates_are_admin_only from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch @@ -95,6 +96,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _update_metadata_fields, _upsert_budget_and_membership, _user_has_admin_view, + validate_budget_duration, ) from litellm.proxy.management_endpoints.organization_endpoints import ( add_member_to_organization, @@ -1153,6 +1155,8 @@ async def new_team( - metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"extra_info": "some info"} - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit for this team - applied across all keys for this team. - model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit for this team - applied across all keys for this team. + - default_estimated_output_tokens: Optional[int] - Expected output tokens reserved for TPM limiting when a request omits max_tokens, for keys on this team that do not set their own. Positive integer. + - default_estimated_output_tokens_per_model: Optional[Dict[str, int]] - Per-model override of the above. Example: {"gpt-4": 4096, "gpt-3.5-turbo": 1024} - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. - tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit @@ -1255,6 +1259,9 @@ async def new_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) + validate_budget_duration(data.budget_duration) + validate_budget_duration(data.team_member_budget_duration) + if data.soft_budget is not None: if data.max_budget is not None: # If max_budget is set, soft_budget must be strictly lower than max_budget @@ -1266,6 +1273,13 @@ async def new_team( }, ) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="team", + ) + # Check if license is over limit total_teams: Final = await _team_db(prisma_client).count() if total_teams and _license_check.is_team_count_over_limit(team_count=total_teams): @@ -1863,6 +1877,8 @@ async def update_team( - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit per model for this team. Example: {"gpt-4": 100, "gpt-3.5-turbo": 200} - model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit per model for this team. Example: {"gpt-4": 10000, "gpt-3.5-turbo": 20000} + - default_estimated_output_tokens: Optional[int] - Expected output tokens reserved for TPM limiting when a request omits max_tokens, for keys on this team that do not set their own. Positive integer. + - default_estimated_output_tokens_per_model: Optional[Dict[str, int]] - Per-model override of the above. Example: {"gpt-4": 4096, "gpt-3.5-turbo": 1024} - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. Example - update team TPM Limit - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. @@ -1935,6 +1951,9 @@ async def update_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) + validate_budget_duration(data.budget_duration) + validate_budget_duration(data.team_member_budget_duration) + existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team_row is None: @@ -1949,6 +1968,14 @@ async def update_team( user_api_key_dict=user_api_key_dict, ) + _existing_team_metadata: Final[object] = getattr(existing_team_row, "metadata", None) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="team", + ) + _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") if data.soft_budget is not None: @@ -2959,7 +2986,7 @@ async def team_member_add( except HTTPException as e: raise e - _validate_budget_duration(data.budget_duration) + validate_budget_duration(data.budget_duration) prisma_client = cast(PrismaClient, prisma_client) @@ -3262,29 +3289,6 @@ def _build_member_budget_patch(data: TeamMemberUpdateRequest) -> dict[str, objec } -def _validate_budget_duration(budget_duration: str | None) -> None: - """Reject budget durations that can't be parsed, are non-positive, or - overflow date math, so a bad value can't be persisted and later crash the - budget reset job.""" - if budget_duration is None: - return - - from litellm.litellm_core_utils.duration_parser import duration_in_seconds - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time - - try: - if duration_in_seconds(budget_duration) <= 0: - raise ValueError("budget_duration must be positive") - get_budget_reset_time(budget_duration=budget_duration) - except (ValueError, OverflowError): - raise HTTPException( - status_code=400, - detail={ - "error": f"Invalid budget_duration '{budget_duration}'. Use a format like '1h', '24h', '7d', or '30d'." - }, - ) - - @router.post( "/team/member_update", tags=["team management"], @@ -3322,7 +3326,7 @@ async def team_member_update( detail={"error": "Either user_id or user_email needs to be passed in"}, ) - _validate_budget_duration(data.budget_duration) + validate_budget_duration(data.budget_duration) _existing_team_row: Final = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py index 987823d987f..3ae8dcf64b7 100644 --- a/litellm/proxy/memory/memory_endpoints.py +++ b/litellm/proxy/memory/memory_endpoints.py @@ -18,13 +18,16 @@ Scoping: """ import json -from typing import Any, Final +from collections.abc import Mapping, Sequence +from datetime import datetime +from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, Depends, HTTPException, Query from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( CommonProxyErrors, + LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth, user_api_key_has_admin_view, @@ -40,10 +43,56 @@ from litellm.types.memory_management import ( MemoryUpdateRequest, ) +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + router: Final = APIRouter() -def _serialize_metadata_for_prisma(metadata: Any) -> str: +class _MemoryRecord(Protocol): + memory_id: str + key: str + value: str + metadata: object + user_id: str | None + team_id: str | None + created_at: datetime | None + created_by: str | None + updated_at: datetime | None + updated_by: str | None + + +class _MemoryTableActions(Protocol): + async def create(self, data: Mapping[str, object]) -> _MemoryRecord: ... + + async def find_many( + self, + where: Mapping[str, object] | None = ..., + order: Mapping[str, str] | None = ..., + skip: int = ..., + take: int = ..., + ) -> Sequence[_MemoryRecord]: ... + + async def count(self, where: Mapping[str, object] | None = ...) -> int: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _MemoryRecord: ... + + async def delete(self, where: Mapping[str, object]) -> _MemoryRecord | None: ... + + +def _memory_table(prisma_client: "PrismaClient") -> _MemoryTableActions: + return MemoryRepository(prisma_client).table + + +class _TeamTableActions(Protocol): + async def find_unique(self, where: Mapping[str, str]) -> LiteLLM_TeamTable | None: ... + + +def _team_table(prisma_client: "PrismaClient") -> _TeamTableActions: + return TeamRepository(prisma_client).table + + +def _serialize_metadata_for_prisma(metadata: object) -> str: """ Encode a `metadata` payload for the `Json?` column. @@ -62,25 +111,25 @@ def _is_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: return user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN -def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> dict | None: +def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object] | None: """ Prisma `where` fragment restricting rows to those the caller can see. Returns None for admins (no restriction). """ if user_api_key_has_admin_view(user_api_key_dict): return None - ors: Final[list[dict]] = [] - if user_api_key_dict.user_id: - ors.append({"user_id": user_api_key_dict.user_id}) - if user_api_key_dict.team_id: - ors.append({"team_id": user_api_key_dict.team_id}) + ors: Final = [ + {field: value} + for field, value in (("user_id", user_api_key_dict.user_id), ("team_id", user_api_key_dict.team_id)) + if value + ] if not ors: # Caller has neither user_id nor team_id — match nothing. return {"memory_id": "__no_match__"} return {"OR": ors} -def _row_to_model(row: Any) -> LiteLLM_MemoryRow: +def _row_to_model(row: _MemoryRecord) -> LiteLLM_MemoryRow: return LiteLLM_MemoryRow( memory_id=row.memory_id, key=row.key, @@ -95,7 +144,7 @@ def _row_to_model(row: Any) -> LiteLLM_MemoryRow: ) -def _require_prisma(): +def _require_prisma() -> "PrismaClient": from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -113,7 +162,9 @@ def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HT return HTTPException(status_code=500, detail=default_detail) -async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: UserAPIKeyAuth) -> None: +async def _assert_write_access( + prisma_client: "PrismaClient", row: _MemoryRecord, user_api_key_dict: UserAPIKeyAuth +) -> None: """ Enforce ownership for mutations (PUT/DELETE). @@ -153,7 +204,7 @@ async def _assert_write_access(prisma_client: Any, row: Any, user_api_key_dict: ) -async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool: +async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: UserAPIKeyAuth, team_id: str) -> bool: """ True if the caller is a team admin of `team_id`, or an org admin for the team's organization. Mirrors the auth pattern used by team-management @@ -168,7 +219,7 @@ async def _is_team_admin_for(prisma_client: Any, user_api_key_dict: UserAPIKeyAu ) try: - team_obj: Final = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + team_obj: Final = await _team_table(prisma_client).find_unique(where={"team_id": team_id}) except Exception as e: verbose_proxy_logger.exception("Error loading team for write-auth check (team_id=%s): %s", team_id, e) return False @@ -269,7 +320,7 @@ async def create_memory( # `metadata` is a `Json?` column — prisma-client-python rejects raw # Python values, so JSON-encode any non-null payload and omit the field # entirely when None so the column defaults to SQL NULL. - create_data: Final[dict] = { + create_data: Final[dict[str, object]] = { "key": body.key, "value": body.value, "user_id": user_id, @@ -281,7 +332,7 @@ async def create_memory( create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) try: - row: Final = await MemoryRepository(prisma_client).table.create(data=create_data) + row: Final = await _memory_table(prisma_client).create(data=create_data) except Exception as e: # Key is globally unique. Any duplicate → 409. if _is_unique_violation(e): @@ -325,14 +376,14 @@ async def list_memory( # top-level "AND" — safer than `dict.update` since future visibility # filters could grow an "OR" key that would clobber this one if merged # by key. - key_filter: Final[dict] = {} + key_filter: Final[dict[str, object]] = {} if key_prefix is not None: key_filter["key"] = {"startsWith": key_prefix} elif key is not None: key_filter["key"] = key vis: Final = _visibility_filter(user_api_key_dict) - where: dict + where: Mapping[str, object] if vis is None: where = key_filter elif not key_filter: @@ -341,8 +392,8 @@ async def list_memory( where = {"AND": [key_filter, vis]} try: - total: Final = await MemoryRepository(prisma_client).table.count(where=where) - rows: Final = await MemoryRepository(prisma_client).table.find_many( + total: Final = await _memory_table(prisma_client).count(where=where) + rows: Final = await _memory_table(prisma_client).find_many( where=where, order={"updated_at": "desc"}, skip=(page - 1) * page_size, @@ -354,12 +405,14 @@ async def list_memory( return MemoryListResponse(memories=[_row_to_model(r) for r in rows], total=total) -async def _find_memory_for_caller(prisma_client: Any, key: str, user_api_key_dict: UserAPIKeyAuth) -> Any: +async def _find_memory_for_caller( + prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth +) -> _MemoryRecord: """Look up a memory row by key, scoped to the caller's visibility.""" - key_filter: Final[dict] = {"key": key} + key_filter: Final[Mapping[str, object]] = {"key": key} vis: Final = _visibility_filter(user_api_key_dict) - where: Final[dict] = key_filter if vis is None else {"AND": [key_filter, vis]} - rows = await MemoryRepository(prisma_client).table.find_many(where=where, take=1, order={"updated_at": "desc"}) + where: Final[Mapping[str, object]] = key_filter if vis is None else {"AND": [key_filter, vis]} + rows = await _memory_table(prisma_client).find_many(where=where, take=1, order={"updated_at": "desc"}) if not rows: raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found") return rows[0] @@ -415,7 +468,7 @@ async def upsert_memory( fields_sent: Final = body.model_fields_set metadata_in_payload: Final = "metadata" in fields_sent - data: Final[dict] = {} + data: Final[dict[str, object]] = {} if body.value is not None: data["value"] = body.value if metadata_in_payload: @@ -427,7 +480,7 @@ async def upsert_memory( ) data["updated_by"] = user_api_key_dict.user_id - async def _find_existing() -> Any: + async def _find_existing() -> _MemoryRecord | None: """Return the caller-visible row for `key`, or None.""" try: return await _find_memory_for_caller(prisma_client, key, user_api_key_dict) @@ -444,7 +497,7 @@ async def upsert_memory( # their team) — otherwise a teammate could overwrite a personal # entry through the OR-based visibility filter. await _assert_write_access(prisma_client, existing, user_api_key_dict) - row = await MemoryRepository(prisma_client).table.update( + row = await _memory_table(prisma_client).update( where={"memory_id": existing.memory_id}, data=data, ) @@ -459,7 +512,7 @@ async def upsert_memory( # Omit `metadata` when None so the column defaults to SQL NULL; # otherwise JSON-encode for Prisma — same pattern as # `create_memory` above. - create_data: Final[dict] = { + create_data: Final[dict[str, object]] = { "key": key, "value": body.value, "user_id": user_id, @@ -470,7 +523,7 @@ async def upsert_memory( if body.metadata is not None: create_data["metadata"] = _serialize_metadata_for_prisma(body.metadata) try: - row = await MemoryRepository(prisma_client).table.create(data=create_data) + row = await _memory_table(prisma_client).create(data=create_data) except Exception as e: # Race: a concurrent PUT/POST created the row after our check. # Re-read and fall back to an update so the PUT stays idempotent @@ -487,7 +540,7 @@ async def upsert_memory( ) # Same write-authorization check as the non-race path. await _assert_write_access(prisma_client, existing_after_race, user_api_key_dict) - row = await MemoryRepository(prisma_client).table.update( + row = await _memory_table(prisma_client).update( where={"memory_id": existing_after_race.memory_id}, data=data, ) @@ -515,7 +568,7 @@ async def delete_memory( # Visibility != write authority — see the upsert handler for the rationale. await _assert_write_access(prisma_client, row, user_api_key_dict) try: - await MemoryRepository(prisma_client).table.delete(where={"memory_id": row.memory_id}) + await _memory_table(prisma_client).delete(where={"memory_id": row.memory_id}) except Exception as e: raise _internal_error("Error deleting memory: %s", e, "Internal error deleting memory entry.") diff --git a/litellm/proxy/openai_files_endpoints/storage_backend_service.py b/litellm/proxy/openai_files_endpoints/storage_backend_service.py index 4c301c96f30..e766f335071 100644 --- a/litellm/proxy/openai_files_endpoints/storage_backend_service.py +++ b/litellm/proxy/openai_files_endpoints/storage_backend_service.py @@ -68,6 +68,16 @@ class StorageBackendFileService: code=400, ) + if target_model_names: + managed_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files") + if not isinstance(managed_files_hook, BaseFileEndpoints): + raise ProxyException( + message="Uploading with target_model_names requires a database-connected proxy, and this proxy has no database configured", + type="invalid_request_error", + param="target_model_names", + code=400, + ) + # Extract file information file_content: Final = file_data["content"] filename: Final = file_data.get("filename", "file") diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 40c49df26cf..f84cdd0c222 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -60,6 +60,7 @@ from .passthrough_endpoint_router import PassthroughEndpointRouter vertex_llm_base: Final = VertexBase() router: Final = APIRouter() +openai_passthrough_router: Final = APIRouter() default_vertex_config: Final = None passthrough_endpoint_router: Final = PassthroughEndpointRouter() @@ -1875,7 +1876,7 @@ async def vertex_proxy_route( ) -@router.api_route( +@openai_passthrough_router.api_route( "/openai_passthrough/{endpoint:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"], tags=["OpenAI Pass-through", "pass-through"], diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 2e3f7bb9aa6..9c5b7dc563e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -1,4 +1,6 @@ +import asyncio import re +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast from urllib.parse import urlparse @@ -39,6 +41,32 @@ else: EndpointType = Any +def _optional_str(value: object) -> str | None: + return value if isinstance(value, str) else None + + +def _optional_str_tuple(value: object) -> tuple[str, ...] | None: + if not isinstance(value, list): + return None + items: Final = cast(list[object], value) # cast-ok: isinstance-narrowed; element type unknown + return tuple(tag for tag in items if isinstance(tag, str)) + + +def _request_tags(request_metadata: Mapping[str, object]) -> tuple[str, ...] | None: + """Tags for the batch-cost spend row: the request's own tags when it sent any, + otherwise the key's tags, which auth exposes as user_api_key_auth_metadata (a + tagged key does not put its tags in the top-level metadata "tags" on the + passthrough path) + """ + tags: Final = _optional_str_tuple(request_metadata.get("tags")) + if tags: + return tags + key_auth_metadata: Final = request_metadata.get("user_api_key_auth_metadata") + if isinstance(key_auth_metadata, dict): + return _optional_str_tuple(key_auth_metadata.get("tags")) + return None + + class VertexPassthroughLoggingHandler: @staticmethod def vertex_passthrough_handler( @@ -657,11 +685,13 @@ class VertexPassthroughLoggingHandler: # Store the managed object for cost tracking # This will be picked up by check_batch_cost polling mechanism + is_batch_create: Final = url_route.split("?")[0].rstrip("/").endswith("batchPredictionJobs") VertexPassthroughLoggingHandler._store_batch_managed_object( unified_object_id=unified_object_id, batch_object=litellm_batch_response, model_object_id=batch_id, logging_obj=logging_obj, + is_batch_create=is_batch_create, **kwargs, ) @@ -779,17 +809,45 @@ class VertexPassthroughLoggingHandler: "kwargs": kwargs, } + @staticmethod + def _log_batch_registration_result( + finished: asyncio.Task, unified_object_id: str, model_object_id: str, is_batch_create: bool + ) -> None: + error: Final = finished.exception() if not finished.cancelled() else None + if finished.cancelled() or error is not None: + consequence: Final = ( + "its cost will not be tracked" if is_batch_create else "its status and output file may be stale" + ) + verbose_proxy_logger.error( + "Failed to store batch managed object with unified_object_id=%s, batch_id=%s; %s: %s", + unified_object_id, + model_object_id, + consequence, + error, + ) + return + verbose_proxy_logger.info( + "Stored batch managed object with unified_object_id=%s, batch_id=%s", + unified_object_id, + model_object_id, + ) + @staticmethod def _store_batch_managed_object( unified_object_id: str, batch_object: LiteLLMBatch, model_object_id: str, logging_obj: LiteLLMLoggingObj, + is_batch_create: bool, **kwargs, ) -> None: """ Store batch managed object for cost tracking. This will be picked up by the check_batch_cost polling mechanism. + + A poll refreshes the batch status and file object but neither creates the row + nor writes attribution, so the creating key and its tags are persisted from + the create alone. """ try: # Get the managed files hook from the logging object @@ -805,7 +863,7 @@ class VertexPassthroughLoggingHandler: user_api_key_dict: Final = UserAPIKeyAuth( user_id=_request_metadata.get("user_api_key_user_id", "default-user"), - api_key="", + api_key=_optional_str(_request_metadata.get("user_api_key")), team_id=_request_metadata.get("user_api_key_team_id"), team_alias=None, user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value @@ -827,9 +885,7 @@ class VertexPassthroughLoggingHandler: ) # Store the unified object for batch cost tracking - import asyncio - - asyncio.create_task( + task: Final = asyncio.create_task( managed_files_hook.store_unified_object_id( unified_object_id=unified_object_id, file_object=batch_object, @@ -837,13 +893,15 @@ class VertexPassthroughLoggingHandler: model_object_id=model_object_id, file_purpose="batch", user_api_key_dict=user_api_key_dict, + request_tags=_request_tags(_request_metadata), + persist_attribution=is_batch_create, + create_if_missing=is_batch_create, ) ) - - verbose_proxy_logger.info( - "Stored batch managed object with unified_object_id=%s, batch_id=%s", - unified_object_id, - model_object_id, + task.add_done_callback( + lambda finished: VertexPassthroughLoggingHandler._log_batch_registration_result( + finished, unified_object_id, model_object_id, is_batch_create + ) ) else: verbose_proxy_logger.warning( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 90daaaeae6b..baa1579d2c0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -15,7 +15,7 @@ import threading import time import traceback import warnings -from collections.abc import AsyncGenerator, Callable, Mapping +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -23,6 +23,7 @@ from typing import ( Any, Final, Literal, + NamedTuple, Optional, TypedDict, Union, @@ -536,6 +537,7 @@ from litellm.proxy.openai_files_endpoints.files_endpoints import ( set_files_config, ) from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + openai_passthrough_router, passthrough_endpoint_router, vertex_ai_live_websocket_passthrough, ) @@ -611,7 +613,7 @@ from litellm.secret_managers.main import ( normalize_nonempty_secret_str, str_to_bool, ) -from litellm.types.integrations.slack_alerting import SlackAlertingArgs +from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs from litellm.types.llms.anthropic import ( AnthropicMessagesRequest, AnthropicResponse, @@ -6859,9 +6861,19 @@ class ProxyConfig: guardrail_id = guardrail.get("guardrail_id") if guardrail_id: db_guardrail_ids.add(guardrail_id) - IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( - guardrail=cast(Guardrail, guardrail), - ) + try: + IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( + guardrail=cast(Guardrail, guardrail), + ) + except Exception as e: # noqa: BLE001 # one unloadable row must not stop the remaining guardrails + verbose_proxy_logger.error( + "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - " + "skipping guardrail '%s' (ID: %s): %s: %s", + guardrail.get("guardrail_name"), + guardrail_id, + type(e).__name__, + e, + ) # Drop in-memory DB-backed entries whose row was deleted on another # pod. Config-loaded entries are never touched. @@ -7632,6 +7644,200 @@ def _pop_complete_sse_frame(buffer: str) -> tuple[str | None, str]: return buffer[:frame_end], buffer[frame_end:] +_STREAM_KEEPALIVE: Final = object() + +_KEEPALIVE_MIN_SECONDS: Final = 1.0 +_KEEPALIVE_MAX_SECONDS: Final = 300.0 +_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) + + +async def _iter_with_keepalive( + aiter: AsyncIterator[Any], + resolve_keepalive_seconds: Callable[[object], float], + keepalive_seconds: float, +) -> AsyncGenerator[Any, None]: + """Wrap `aiter` with idle-gap heartbeats, re-resolving the interval after each + real chunk via `resolve_keepalive_seconds`. A mid-stream router fallback can + swap in a deployment with a different keepalive policy, including one that + newly enables or newly disables heartbeats, partway through the same stream; + re-resolving against each chunk's own identity (rather than trusting the + interval picked before iteration started, or picked the last time it went + inactive) keeps the heartbeat behavior in sync with whichever deployment + actually produced it, in both directions. While the interval is <= 0, no + task is created and no timeout is awaited: a chunk is forwarded the moment + it arrives, at the same cost as a bare `async for`.""" + pending: asyncio.Task[Any] | None = None # rebind-ok: rebound each loop iteration + current_keepalive_seconds = keepalive_seconds # rebind-ok: re-resolved after each chunk + try: + while True: + if current_keepalive_seconds <= 0: + try: + item = await aiter.__anext__() + except StopAsyncIteration: + break + yield item + current_keepalive_seconds = resolve_keepalive_seconds(item) + continue + + if pending is None: + pending = asyncio.create_task(aiter.__anext__()) + done, _ = await asyncio.wait((pending,), timeout=current_keepalive_seconds) + if not done: + yield _STREAM_KEEPALIVE + continue + try: + item = pending.result() + except StopAsyncIteration: + break + finally: + pending = None + yield item + current_keepalive_seconds = resolve_keepalive_seconds(item) + finally: + if pending is not None and not pending.done(): + pending.cancel() + try: + await pending + except asyncio.CancelledError: + pass + + +class _DeploymentKeepaliveConfig(NamedTuple): + keepalive_seconds: Any + allow_client_override: bool + + +def _keepalive_from_deployment_config( + request_data: Mapping[str, Any], response: object +) -> _DeploymentKeepaliveConfig | None: + if llm_router is None: + return None + + hidden: Final = get_hidden_params_dict(response) + model_id: Final = hidden.get("model_id") + if isinstance(model_id, str) and model_id: + deployment: Final = llm_router.get_deployment(model_id=model_id) + # A populated model_id names the specific deployment that served this + # stream. If it no longer resolves (e.g. removed by a config reload + # mid-stream), that's a stale identity, not an absent one: don't fall + # through to guessing via model_name below, since a currently-live + # sibling deployment's config was never what actually served this + # stream. + if deployment is None: + return None + return _DeploymentKeepaliveConfig( + keepalive_seconds=getattr(deployment.litellm_params, "keepalive_seconds", None), + allow_client_override=bool(getattr(deployment.litellm_params, "allow_client_keepalive_override", False)), + ) + + # No model_id at all to pin down which deployment actually served this + # stream: only trust the fallback when every deployment under this + # model_name agrees on both keepalive_seconds and + # allow_client_keepalive_override (including deployments that leave either + # field unset), so a stream never inherits a sibling deployment's policy. + configs: Final = frozenset( + ( + (deployment_dict.get("litellm_params") or _EMPTY_MAPPING).get("keepalive_seconds"), + bool( + (deployment_dict.get("litellm_params") or _EMPTY_MAPPING).get("allow_client_keepalive_override", False) + ), + ) + for deployment_dict in llm_router.get_model_list(model_name=request_data.get("model")) or () + ) + if len(configs) == 1: + keepalive_seconds, allow_client_override = next(iter(configs)) + return _DeploymentKeepaliveConfig( + keepalive_seconds=keepalive_seconds, allow_client_override=allow_client_override + ) + return None + + +def _is_explicit_keepalive_disable(raw: object) -> bool: + if not isinstance(raw, (int, float, str)): + return False + try: + return float(raw) <= 0 + except ValueError: + return False + + +def _resolve_keepalive_seconds(request_data: Mapping[str, Any], response: object = None) -> float: + deployment_config: Final = _keepalive_from_deployment_config(request_data, response) + deployment_raw: Final = deployment_config.keepalive_seconds if deployment_config is not None else None + allow_client_override: Final = deployment_config.allow_client_override if deployment_config is not None else False + + # An operator setting keepalive_seconds: 0 on a deployment is an explicit hard + # disable: an authenticated client must not be able to re-enable heartbeats + # (and the idle-timeout evasion that comes with them) for a deployment the + # operator opted out of, regardless of what the request body asks for. + if _is_explicit_keepalive_disable(deployment_raw): + return 0.0 + + # keepalive_seconds is operator-only unless the deployment explicitly opts in: + # a client can't unilaterally enable heartbeats (and the LB-idle-timeout + # evasion that comes with them) for a deployment that never configured this. + client_supplied: Final = request_data.get("keepalive_seconds") if allow_client_override else None + raw: Final = client_supplied if client_supplied is not None else deployment_raw + try: + value: Final = float(raw) if isinstance(raw, (int, float, str)) else 0.0 + except ValueError: + return 0.0 + if value <= 0: + return 0.0 + clamped: Final = max(_KEEPALIVE_MIN_SECONDS, min(value, _KEEPALIVE_MAX_SECONDS)) + if clamped != value: + verbose_proxy_logger.info( + "keepalive_seconds=%s clamped to %s [min=%s, max=%s]", + value, + clamped, + _KEEPALIVE_MIN_SECONDS, + _KEEPALIVE_MAX_SECONDS, + ) + return clamped + + +_KEEPALIVE_CACHE_TTL_SECONDS: Final = 5.0 + + +def _make_keepalive_resolver(request_data: Mapping[str, Any]) -> Callable[[object], float]: + """Wrap `_resolve_keepalive_seconds` with a memo keyed on the serving + deployment's model_id. The steady-state case (no mid-stream fallback, the + overwhelming majority of streams) sees the same model_id on every chunk, so + this turns the per-chunk cost from a full `llm_router.get_deployment()` + Pydantic rebuild into a cheap hidden-params read once per + `_KEEPALIVE_CACHE_TTL_SECONDS` for that model_id. The cache expires on its + own rather than living for the life of the stream, so an operator's live + config change (disabling keepalive, revoking client override, or removing + the deployment) is observed within a bounded window instead of being able + to be evaded by an already-in-flight stream indefinitely. A missing/empty + model_id can't be trusted as a cache key (see + `_keepalive_from_deployment_config`'s model_name fallback, which reflects + current router state rather than one deployment's fixed identity), so + those chunks always resolve fresh, matching prior behavior exactly. + """ + last_model_id: str | None = None # rebind-ok: memoized identity of the last-resolved chunk + last_value: float = 0.0 # rebind-ok: cached resolution for last_model_id + last_resolved_at: float = float("-inf") # rebind-ok: monotonic timestamp of the last real resolution + + def _resolve(item: object) -> float: + nonlocal last_model_id, last_value, last_resolved_at + model_id = get_hidden_params_dict(item).get("model_id") + now: Final = time.monotonic() + if ( + isinstance(model_id, str) + and model_id + and model_id == last_model_id + and now - last_resolved_at < _KEEPALIVE_CACHE_TTL_SECONDS + ): + return last_value + value: Final = _resolve_keepalive_seconds(request_data, item) + if isinstance(model_id, str) and model_id: + last_model_id, last_value, last_resolved_at = model_id, value, now + return value + + return _resolve + + async def async_data_generator( response, user_api_key_dict: UserAPIKeyAuth, @@ -7680,7 +7886,28 @@ async def async_data_generator( else: stream_iterator = response - async for chunk in stream_iterator: + # A stream can start on a deployment with keepalive off and fall back + # mid-stream to one that enables it: only skip wrapping altogether when + # there's no router to ever fall back through in the first place (in + # which case _resolve_keepalive_seconds can never return non-zero for + # any chunk of this stream), not merely because the first chunk's + # deployment happens to start with it off. + resolve_keepalive_seconds: Final = _make_keepalive_resolver(request_data) + stream_source: Final = ( + _iter_with_keepalive( + stream_iterator.__aiter__(), + resolve_keepalive_seconds, + resolve_keepalive_seconds(response), + ) + if llm_router is not None + else stream_iterator + ) + + async for item in stream_source: + if item is _STREAM_KEEPALIVE: + yield ": ping\n\n" + continue + chunk = cast(Any, item) # cast-ok: sentinel already handled above, item is a real chunk here if needs_per_chunk_hook: ### CALL HOOKS ### - modify outgoing data chunk, _str_so_far = await _apply_streaming_chunk_hooks( @@ -8469,6 +8696,47 @@ class ProxyStartupEvent: await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) + ### PTU DAILY ROLLUP ### + from litellm.proxy.spend_tracking.ptu_feature_flag import ( + is_ptu_cost_attribution_enabled, + ) + + if is_ptu_cost_attribution_enabled(): + from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( + PTU_ROLLUP_JOB_ID, + run_scheduled_ptu_rollup, + ) + + async def _alert_ptu_rollup_failure(message: str) -> None: + await proxy_logging_obj.alerting_handler( + message=message, + level="High", + alert_type=AlertType.failed_tracking_spend, + ) + + async def _scheduled_ptu_rollup() -> None: + # Reuse the PodLockManager from db_spend_update_writer so only one pod + # reconciles a day; a multi-pod race could prune another pod's fresh rows + await run_scheduled_ptu_rollup( + prisma_client, + pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, + alert=_alert_ptu_rollup_failure, + ) + + scheduler.add_job( + _scheduled_ptu_rollup, + "cron", + hour=0, + minute=15, + timezone="UTC", + id=PTU_ROLLUP_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info( + "PTU rollup job scheduled at 00:15 UTC daily (only models with PTU config accrue flat cost)" + ) + ### SPEND LOG CLEANUP ### if ( general_settings.get("maximum_spend_logs_retention_period") is not None @@ -16607,6 +16875,7 @@ app.include_router(search_router) app.include_router(image_router) app.include_router(fine_tuning_router) app.include_router(credential_router) +app.include_router(openai_passthrough_router) app.include_router(batches_router) app.include_router(openai_files_router) app.include_router(llm_passthrough_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9c871b65f40..33fd9389b63 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -30,7 +30,7 @@ model LiteLLM_BudgetTable { end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team - organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } // Models on proxy @@ -893,6 +893,7 @@ model LiteLLM_DailyTeamSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) created_at DateTime @default(now()) updated_at DateTime @updatedAt @@ -985,6 +986,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + api_key String? + request_tags Json? @default("[]") updated_at DateTime @updatedAt updated_by String? diff --git a/litellm/proxy/spend_tracking/ptu_feature_flag.py b/litellm/proxy/spend_tracking/ptu_feature_flag.py new file mode 100644 index 00000000000..9078079b676 --- /dev/null +++ b/litellm/proxy/spend_tracking/ptu_feature_flag.py @@ -0,0 +1,18 @@ +"""Opt-in flag for PTU (provisioned throughput unit) flat-cost attribution. + +The whole feature is inert unless an operator sets +``LITELLM_ENABLE_PTU_COST_ATTRIBUTION``: the daily rollup is not scheduled, the +model endpoints reject PTU config, the daily activity read path reports zero flat +cost, and the model form hides the PTU inputs. +""" + +from typing import Final + +from litellm.secret_managers.main import get_secret_bool + +PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION" + + +def is_ptu_cost_attribution_enabled() -> bool: + """Report whether this deployment opted into PTU flat-cost attribution.""" + return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py new file mode 100644 index 00000000000..029648f7901 --- /dev/null +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -0,0 +1,663 @@ +""" +Daily rollup for per-model PTU (provisioned throughput) flat cost. + +v1 reads PTU config straight off the model deployment +(``LiteLLM_ProxyModelTable.model_info``): a deployment carrying ``ptu_count`` +and ``cost_per_ptu_per_hour`` accrues flat cost of +``ptu_count * cost_per_ptu_per_hour * active_hours`` for a given UTC day, where +``active_hours`` is the overlap between the day and the optional +``[ptu_effective_from, ptu_effective_to)`` window (a window opening at 23:00 +charges one hour that day). The amount is written to ``LiteLLM_DailyTeamSpend`` +under a sentinel api_key so the rows are distinguishable from per-request rows +and share the existing unique constraint. +""" + +import asyncio +import json +from collections.abc import Awaitable, Callable, Mapping +from dataclasses import dataclass +from datetime import date, datetime, time, timedelta, timezone +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + PTU_PRUNE_SKEW_GRACE_SECONDS, + PTU_ROLLUP_JOB_ID, + PTU_ROLLUP_LOCK_TTL_SECONDS, + PTU_ROLLUP_MAX_BACKFILL_DAYS, + PTU_SENTINEL_API_KEY, +) +from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled +from litellm.types.router import ModelInfo + +if TYPE_CHECKING: + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + from litellm.proxy.utils import PrismaClient + +_HOURS_PER_DAY: Final = 24 +_UPSERT_ATTEMPTS: Final = 3 +_UPSERT_RETRY_BACKOFF_SECONDS: Final = 0.5 + + +@dataclass(frozen=True, slots=True) +class RollupResult: + day: date + models_processed: int + rows_written: int + rows_failed: int = 0 + + +@dataclass(frozen=True, slots=True) +class BackfillResult: + start: date + end: date + days_scanned: int + rows_written: int + rows_failed: int = 0 + + +@dataclass(frozen=True, slots=True) +class PTUModel: + """A model deployment carrying valid manual PTU config.""" + + model_id: str + model_name: str + team_id: str + ptu_count: int + cost_per_ptu_per_hour: float + effective_from: datetime | None = None + effective_to: datetime | None = None + + +def _parse_utc_datetime(value: object) -> datetime | None: + """Parse a model_info datetime (ISO string or datetime) into a UTC-aware datetime, else None.""" + parsed: Final = _coerce_datetime(value) + if parsed is None: + return None + if parsed.tzinfo is None: + return parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _coerce_datetime(value: object) -> datetime | None: + """``value`` as a datetime, parsing an ISO string, else None.""" + if isinstance(value, datetime): + return value + if not isinstance(value, str): + return None + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")) + except ValueError: + return None + + +def _public_model_name(row: object, model_info: Mapping[str, object]) -> str: + """The name an operator recognises for this deployment. + + Creating a team-scoped deployment rewrites model_name to a synthetic routing key + (``model_name__``) and keeps the chosen name in + ``model_info.team_public_model_name``. PTU config is only accepted alongside a + team_id, so every PTU deployment carries that synthetic name; keying the sentinel + row on it would file each charge under a UUID that no usage view can resolve and + that never lines up with the same model's request rows. + """ + public_name: Final = model_info.get("team_public_model_name") + if isinstance(public_name, str) and public_name: + return public_name + return str(getattr(row, "model_name", "") or "") + + +def _decode_model_info(raw: object) -> "Mapping[str, object] | None": + """A deployment's model_info as a dict, decoding a JSON string, else None.""" + if isinstance(raw, str): + try: + return json.loads(raw) + except (TypeError, ValueError): + return None + if isinstance(raw, dict): + return raw + return None + + +def _parse_ptu_model(row: object) -> PTUModel | None: + """Return a PTUModel when the deployment carries valid manual PTU config, else None. + + Valid means model_info has a positive ptu_count, a non-negative + cost_per_ptu_per_hour, and a team_id (1 model -> 1 team). + """ + raw_model_info: Final = getattr(row, "model_info", None) + model_info: Final = _decode_model_info(raw_model_info) + if model_info is None: + return None + ptu_count: Final = model_info.get("ptu_count") + cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour") + team_id: Final = model_info.get("team_id") + if ptu_count is None or cost_per_hour is None or not team_id: + return None + try: + ptu_count_int: Final = int(ptu_count) + cost_per_hour_float: Final = float(cost_per_hour) + except (TypeError, ValueError, OverflowError): + return None + if not 0 < ptu_count_int <= ModelInfo.MAX_PTU_COUNT: + return None + if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR: + return None + if model_info.get("ptu_effective_from") is None: + # The endpoints require a start; a row without one predates that rule or was + # written around them, and inferring one would bill days the deployment did not exist + return None + raw_from: Final = model_info.get("ptu_effective_from") + raw_to: Final = model_info.get("ptu_effective_to") + effective_from: Final = _parse_utc_datetime(raw_from) + effective_to: Final = _parse_utc_datetime(raw_to) + # A present-but-unparseable bound would read as "no bound" and silently widen the + # window to the whole day, so the deployment is skipped until the config is fixed + if (raw_from is not None and effective_from is None) or (raw_to is not None and effective_to is None): + return None + if effective_from is not None and effective_to is not None and effective_to <= effective_from: + return None + return PTUModel( + model_id=str(getattr(row, "model_id", "") or ""), + model_name=_public_model_name(row, model_info), + team_id=str(team_id), + ptu_count=ptu_count_int, + cost_per_ptu_per_hour=cost_per_hour_float, + effective_from=effective_from, + effective_to=effective_to, + ) + + +def _active_hours_on_day(model: PTUModel, day: date) -> float: + """Hours the model's PTU window overlaps ``day`` (UTC), clamped to [0, 24].""" + day_start: Final = datetime.combine(day, time.min, tzinfo=timezone.utc) + day_end: Final = day_start + timedelta(days=1) + start: Final = max(day_start, model.effective_from) if model.effective_from else day_start + end: Final = min(day_end, model.effective_to) if model.effective_to else day_end + if end <= start: + return 0.0 + return (end - start).total_seconds() / 3600.0 + + +def _compute_daily_flat_cost(model: PTUModel, day: date) -> float: + """Flat cost for ``day``: ptu_count * cost_per_ptu_per_hour * active_hours.""" + return float(model.ptu_count) * model.cost_per_ptu_per_hour * _active_hours_on_day(model, day) + + +@dataclass(frozen=True, slots=True) +class _PTUCharge: + """One sentinel row's worth of flat cost for a deployment on a day. + + ``model_id`` is the row's identity and goes in the unique key; ``model_name`` is what + an operator reads and rides alongside it. A deployment can be renamed, so keying on + the name would let two runs holding different config views write the same day twice. + """ + + team_id: str + model_id: str + model_name: str + flat_cost: float + + +def _aggregate_charges(ptu_models: tuple[PTUModel, ...], day: date) -> tuple[_PTUCharge, ...]: + """One charge per deployment that accrues cost on ``day``. Zero-cost deployments are + dropped, which keeps a day outside a window from writing a row. + + Deployments sharing a public name inside a team no longer need collapsing: each keys + its own row on its own id, and the read path merges them back under the shared name. + """ + return tuple( + _PTUCharge( + team_id=model.team_id, + model_id=model.model_id, + model_name=model.model_name, + flat_cost=_compute_daily_flat_cost(model, day), + ) + for model in sorted(ptu_models, key=lambda m: (m.team_id, m.model_id)) + if _compute_daily_flat_cost(model, day) > 0 + ) + + +async def _upsert_ptu_daily_row( + prisma_client: "PrismaClient", + *, + team_id: str, + model_id: str, + model_name: str, + date_str: str, + flat_cost: float, +) -> None: + """Idempotent upsert of a sentinel-api_key row on LiteLLM_DailyTeamSpend. + + ``model`` holds the deployment id because it is part of the table's unique key and a + rename must not move the row. ``model_group`` carries the operator-facing name, which + is outside the key and is what the usage views display. + """ + where: Final = { # mutable-ok: prisma upsert filter payload + "team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": { # mutable-ok: prisma composite-key filter + "team_id": team_id, + "date": date_str, + "api_key": PTU_SENTINEL_API_KEY, + "model": model_id, + "custom_llm_provider": "", + "mcp_namespaced_tool_name": "", + "endpoint": "", + } + } + now: Final = datetime.now(timezone.utc) + await prisma_client.db.litellm_dailyteamspend.upsert( + where=where, + data={ # mutable-ok: prisma upsert data payload + "create": { # mutable-ok: prisma create payload + "team_id": team_id, + "date": date_str, + "api_key": PTU_SENTINEL_API_KEY, + "model": model_id, + "model_group": model_name, + "custom_llm_provider": "", + "mcp_namespaced_tool_name": "", + "endpoint": "", + "ptu_flat_cost": flat_cost, + }, + "update": { # mutable-ok: prisma update payload + "model_group": model_name, + "ptu_flat_cost": flat_cost, + "updated_at": now, + }, + }, + ) + + +async def _upsert_charge_with_retry( + prisma_client: "PrismaClient", + *, + charge: _PTUCharge, + date_str: str, +) -> bool: + """Write one charge, retrying transient failures. Returns False once attempts are spent. + + The upsert is idempotent on the sentinel unique key, so a retry can only rewrite the + same amount for the same day. Retrying in-run matters because the scheduled job moves + on to the next date: a write lost here is a day of PTU cost that no later run replays. + """ + for attempt in range(1, _UPSERT_ATTEMPTS + 1): + try: + await _upsert_ptu_daily_row( + prisma_client, + team_id=charge.team_id, + model_id=charge.model_id, + model_name=charge.model_name, + date_str=date_str, + flat_cost=charge.flat_cost, + ) + return True + except Exception as exc: # noqa: BLE001 # one bad row must not stop the batch + if attempt < _UPSERT_ATTEMPTS: + verbose_proxy_logger.warning( + "PTU rollup: upsert attempt %d/%d failed for team=%s model=%s day=%s: %s", + attempt, + _UPSERT_ATTEMPTS, + charge.team_id, + charge.model_name, + date_str, + exc, + ) + await asyncio.sleep(_UPSERT_RETRY_BACKOFF_SECONDS * attempt) + continue + verbose_proxy_logger.error( + "PTU rollup: upsert failed after %d attempts for team=%s model=%s day=%s " + "(rerun the rollup for that date to recover): %s", + _UPSERT_ATTEMPTS, + charge.team_id, + charge.model_name, + date_str, + exc, + ) + return False + + +async def _load_ptu_models(prisma_client: "PrismaClient") -> tuple[PTUModel, ...]: + """Every model deployment currently carrying valid manual PTU config.""" + rows: Final = await prisma_client.db.litellm_proxymodeltable.find_many() + return tuple(parsed for parsed in (_parse_ptu_model(row) for row in rows) if parsed is not None) + + +async def run_ptu_flat_cost_rollup( + prisma_client: "PrismaClient", + target_date: date | None = None, + may_prune: bool = True, +) -> RollupResult: + """Rollup one UTC day of flat PTU cost across all PTU-configured model deployments. + + Defaults to yesterday UTC. Authoritative for the day: it upserts the current charges + first, then deletes the day's sentinel rows this run did not refresh, so a + since-removed, invalidated, or now-out-of-window deployment leaves no stale charge. + + The prune predicate is ``updated_at < run_started`` rather than "not in the charge + set I computed", which matters under concurrency: whether a row is garbage becomes a + property of the row instead of one run's in-memory config snapshot, so a run can + never delete a row a concurrent run just wrote. It is still skipped when any charge + failed to write, since a row whose replacement never landed would look unrefreshed. + """ + day: Final = target_date or (datetime.now(timezone.utc).date() - timedelta(days=1)) + + if prisma_client is None: + verbose_proxy_logger.warning("PTU rollup: prisma_client is None, skipping") + return RollupResult(day=day, models_processed=0, rows_written=0) + + date_str: Final = day.isoformat() + run_started: Final = datetime.now(timezone.utc) + + ptu_models: Final = await _load_ptu_models(prisma_client) + charges: Final = _aggregate_charges(ptu_models, day) + + landed: Final = tuple( + [await _upsert_charge_with_retry(prisma_client, charge=charge, date_str=date_str) for charge in charges] + ) + rows_written: Final = sum(landed) + rows_failed: Final = len(charges) - rows_written + + if not may_prune: + verbose_proxy_logger.info( + "PTU rollup for %s: ran without the cross-pod lock, skipping the prune so a " + "concurrent pod's charges cannot be swept by this run's cutoff", + date_str, + ) + elif rows_failed: + # A charge that never landed leaves its row looking unrefreshed, so the prune + # would delete the very row the failed write was meant to replace + verbose_proxy_logger.warning( + "PTU rollup: %d charge(s) failed for %s, skipping the prune so a row whose " + "replacement did not land is not deleted; rerun that date to reconcile", + rows_failed, + date_str, + ) + else: + await _prune_unrefreshed_sentinel_rows(prisma_client, date_str=date_str, run_started=run_started) + + verbose_proxy_logger.info( + "PTU rollup for %s: %d PTU models processed, %d rows written, %d rows failed", + date_str, + len(ptu_models), + rows_written, + rows_failed, + ) + return RollupResult( + day=day, + models_processed=len(ptu_models), + rows_written=rows_written, + rows_failed=rows_failed, + ) + + +def _backfill_window(ptu_models: tuple[PTUModel, ...], end: date) -> tuple[date, ...]: + """The UTC days the catch-up pass considers, oldest first, through ``end`` inclusive. + + Starts at the earliest declared ``ptu_effective_from``, floored at + ``PTU_ROLLUP_MAX_BACKFILL_DAYS`` before ``end``. A start is required alongside the + count and rate, so a deployment without one is not priced rather than being given the + floor, which would bill it for the whole cap window. Empty when there is no PTU + config, or when every declared window opens after ``end``. + """ + floor: Final = end - timedelta(days=PTU_ROLLUP_MAX_BACKFILL_DAYS) + starts: Final = tuple(model.effective_from.date() for model in ptu_models if model.effective_from) + if not starts: + return () + start: Final = max(min(starts), floor) + return tuple(start + timedelta(days=offset) for offset in range((end - start).days + 1)) + + +async def _existing_sentinel_keys( + prisma_client: "PrismaClient", + *, + start: date, + end: date, +) -> frozenset[tuple[str, str, str]]: + """``(team_id, deployment id, date)`` of every PTU sentinel row within ``[start, end]``. + + The row's ``model`` column holds the deployment id, so this is an exact identity and + survives a rename. Nothing here reads the display name. + """ + date_range: Final = {"gte": start.isoformat(), "lte": end.isoformat()} # mutable-ok: prisma range filter + rows: Final = await prisma_client.db.litellm_dailyteamspend.find_many( + where={"api_key": PTU_SENTINEL_API_KEY, "date": date_range} # mutable-ok: prisma find filter + ) + return frozenset( + ( + str(getattr(row, "team_id", "") or ""), + str(getattr(row, "model", "") or ""), + str(getattr(row, "date", "") or ""), + ) + for row in rows + ) + + +async def run_ptu_flat_cost_backfill( + prisma_client: "PrismaClient", + today: date | None = None, +) -> BackfillResult: + """Price the elapsed days of every PTU window that carry no sentinel row yet. + + Writes only the charges that are missing and never rewrites or deletes an existing + row, so a day already priced keeps the amount it was billed, whatever the config says + now. A day counts as priced when a sentinel row exists for that deployment id, so + renaming a deployment neither re-prices its history nor files a second charge beside + the row already there. Zero-cost days write nothing, which leaves a day + outside a window reconsidered on each run rather than recorded as done. + + It deletes nothing. Removing a deployment stops it accruing new charges and leaves the + days it was billed for standing, since those days were incurred. + """ + end: Final = (today or datetime.now(timezone.utc).date()) - timedelta(days=1) + + if not prisma_client: + verbose_proxy_logger.warning("PTU backfill: prisma_client is None, skipping") + return BackfillResult(start=end, end=end, days_scanned=0, rows_written=0) + + ptu_models: Final = await _load_ptu_models(prisma_client) + days: Final = _backfill_window(ptu_models, end) + + if not days: + return BackfillResult(start=end, end=end, days_scanned=0, rows_written=0) + + priced: Final = await _existing_sentinel_keys(prisma_client, start=days[0], end=days[-1]) + missing: Final = tuple( + (day.isoformat(), charge) + for day in days + for charge in _aggregate_charges(ptu_models, day) + if (charge.team_id, charge.model_id, day.isoformat()) not in priced + ) + if not missing: + return BackfillResult(start=days[0], end=days[-1], days_scanned=len(days), rows_written=0) + + landed: Final = tuple( + [ + await _upsert_charge_with_retry(prisma_client, charge=charge, date_str=date_str) + for date_str, charge in missing + ] + ) + rows_written: Final = sum(landed) + verbose_proxy_logger.info( + "PTU backfill for %s to %s: %d unpriced charge(s) found, %d written, %d failed", + days[0].isoformat(), + days[-1].isoformat(), + len(missing), + rows_written, + len(missing) - rows_written, + ) + return BackfillResult( + start=days[0], + end=days[-1], + days_scanned=len(days), + rows_written=rows_written, + rows_failed=len(missing) - rows_written, + ) + + +async def run_scheduled_ptu_rollup( + prisma_client: "PrismaClient", + pod_lock_manager: "PodLockManager | None" = None, + target_date: date | None = None, + alert: Callable[[str], Awaitable[None]] | None = None, +) -> RollupResult | None: + """Run the daily rollup under a cross-pod lock so only one proxy reconciles a day. + + Every proxy process schedules this cron, and the read-charge-prune sequence is not + atomic: two pods reading different config snapshots can have the loser's prune delete + a row the winner just wrote. Returns None when another pod holds the lock, since that + pod is doing the work. A deployment without a Redis-backed lock manager runs + unguarded, as ``SpendLogCleanup`` does, and so does a run that cannot reach Redis at + all: the lock exists to avoid duplicate work, so no lock problem may cost a day. + + The lease is a fixed TTL with no renewal, so a long scan can outlive it. That costs + duplicate work rather than correctness: the upserts are idempotent on the sentinel + key and the prune reads only the row's own timestamp, so a second pod arriving + mid-run cannot corrupt the day. + + Returns None without touching the database when PTU cost attribution is off. Proxy + startup already skips scheduling the cron, so this guards the function itself rather + than its one caller, and a deployment that never opted in accrues nothing whatever + reaches it. + """ + if not is_ptu_cost_attribution_enabled(): + return None + + if pod_lock_manager is None or pod_lock_manager.redis_cache is None: + return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False) + + if not await pod_lock_manager.acquire_lock(cronjob_id=PTU_ROLLUP_JOB_ID, ttl=PTU_ROLLUP_LOCK_TTL_SECONDS): + if await _lock_is_held(pod_lock_manager): + verbose_proxy_logger.info("PTU rollup: another pod holds the rollup lock, skipping this run") + return None + # acquire_lock reports contention and a Redis outage the same way, so an + # unreachable Redis would otherwise skip the day on every pod at once. The + # reconcile is safe to run concurrently, so losing the lock costs duplicate + # work; losing the day costs a team's charges + verbose_proxy_logger.warning( + "PTU rollup: could not take the rollup lock and no other pod holds it, " + "running unguarded rather than skipping the day" + ) + return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False) + + try: + return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=True) + finally: + await pod_lock_manager.release_lock(cronjob_id=PTU_ROLLUP_JOB_ID) + + +async def _lock_is_held(pod_lock_manager: "PodLockManager") -> bool: + """True only when the rollup lock is readable and someone is holding it. + + A Redis that cannot be read is reported as "not held" so the caller runs the day + rather than skipping it; the cost of being wrong here is a duplicate reconcile. + """ + try: + lock_key: Final = pod_lock_manager.get_redis_lock_key(PTU_ROLLUP_JOB_ID) + return bool(await pod_lock_manager.redis_cache.async_get_cache(lock_key)) + except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the day + verbose_proxy_logger.warning("PTU rollup: could not read the rollup lock: %s", exc) + return False + + +async def _run_and_alert( + prisma_client: "PrismaClient", + *, + target_date: date | None, + alert: "Callable[[str], Awaitable[None]] | None", + may_prune: bool = True, +) -> RollupResult: + """Reconcile the day, catch up any days left unpriced, and alert on charges that did not land. + + A charge that exhausts its retries leaves that team showing no PTU cost for the date, + and the scheduled job moves on to the next day rather than replaying it. That is a + silent underbill unless someone is reading proxy logs, so it is escalated to whatever + alerting the deployment has configured. + + The catch-up pass runs only on the scheduled shape, where ``target_date`` is None. An + explicit date means reconcile exactly that day, so it stays a single-day operation. + Its failure is contained: the day's own result is returned either way. + """ + result: Final = await run_ptu_flat_cost_rollup(prisma_client, target_date=target_date, may_prune=may_prune) + if result.rows_failed: + await _deliver_alert( + alert, + f"PTU flat-cost rollup for {result.day.isoformat()}: {result.rows_failed} of " + f"{result.rows_written + result.rows_failed} team charges failed to write. Those teams show no PTU " + f"cost for that date until the rollup is rerun for it.", + ) + if target_date is None: + await _backfill_and_alert(prisma_client, alert=alert) + return result + + +async def _backfill_and_alert( + prisma_client: "PrismaClient", + *, + alert: "Callable[[str], Awaitable[None]] | None", +) -> None: + """Catch up unpriced PTU days, alerting on charges that did not land. + + Never raises: the day's own rollup has already run and its result must reach the + caller whatever the catch-up pass does. + """ + try: + backfill: Final = await run_ptu_flat_cost_backfill(prisma_client) + except Exception as exc: # noqa: BLE001 # the catch-up pass must not fail the day's rollup + verbose_proxy_logger.error("PTU backfill: catch-up pass failed, the day's rollup still stands: %s", exc) + return + if backfill.rows_failed: + await _deliver_alert( + alert, + f"PTU flat-cost backfill for {backfill.start.isoformat()} to {backfill.end.isoformat()}: " + f"{backfill.rows_failed} of {backfill.rows_written + backfill.rows_failed} previously unpriced charges " + f"failed to write. Those days stay unpriced until a later run picks them up.", + ) + + +async def _deliver_alert(alert: "Callable[[str], Awaitable[None]] | None", message: str) -> None: + """Send an operator alert when one is configured, swallowing a broken channel.""" + if alert is None: + return + try: + await alert(message) + except Exception as exc: # noqa: BLE001 # a broken alert channel must not fail the rollup + verbose_proxy_logger.error("PTU rollup: could not deliver the failed-charge alert: %s", exc) + + +async def _prune_unrefreshed_sentinel_rows( + prisma_client: "PrismaClient", + *, + date_str: str, + run_started: datetime, +) -> None: + """Delete the day's PTU sentinel rows this run did not refresh. + + Every charge the run wrote bumps ``updated_at`` past ``run_started``, so anything + left below that mark is a (team, model) the current config no longer prices. The mark + is pulled back by ``PTU_PRUNE_SKEW_GRACE_SECONDS`` because the two timestamps come + from different hosts: a stale row is hours old, a concurrently written one is seconds + old, and the grace separates them without waiting on clocks agreeing. The + predicate reads only the row, never the caller's config snapshot, which is what + makes it safe to run twice, out of order, or beside another pod: a row written + after this run began is out of reach of its delete. Mirrors the retention predicate + ``SpendLogCleanup`` deletes by.""" + cutoff: Final = run_started - timedelta(seconds=PTU_PRUNE_SKEW_GRACE_SECONDS) + await prisma_client.db.litellm_dailyteamspend.delete_many( + where={ # mutable-ok: prisma delete filter + "date": date_str, + "api_key": PTU_SENTINEL_API_KEY, + "updated_at": {"lt": cutoff}, # mutable-ok: prisma comparison filter + } + ) + + +__all__ = ( + "PTU_ROLLUP_JOB_ID", + "PTU_SENTINEL_API_KEY", + "BackfillResult", + "PTUModel", + "RollupResult", + "run_ptu_flat_cost_backfill", + "run_ptu_flat_cost_rollup", + "run_scheduled_ptu_rollup", +) diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index 3332afc0a4b..448723ab3bc 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Final, NamedTuple import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.llm_cost_calc.utils import _get_cost_per_unit, generic_cost_per_token if TYPE_CHECKING: from litellm.router import Router @@ -26,29 +26,42 @@ class SavingsSpend(NamedTuple): autorouter: float = 0.0 -def _input_and_cache_read_cost(model: str | None, custom_llm_provider: str | None) -> tuple[float, float]: +def _input_cache_read_and_write_cost(info: ModelInfo | None) -> tuple[float, float, float]: """ - Return ``(input_cost_per_token, cache_read_cost_per_token)`` for a model. + Return ``(input_cost, cache_read_cost, cache_write_cost)`` per token. - Falls open to ``(0.0, 0.0)`` when the model is unknown so savings degrade to - zero rather than raising inside the spend writer. When a model has no - separate cache-read price the cache-read cost mirrors the input cost, which - yields zero caching savings. + ``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 not model: - return 0.0, 0.0 - try: - info: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; degrade to zero savings - verbose_proxy_logger.debug( - "savings: no model info for provider=%s model=%s (%s)", custom_llm_provider, model, e - ) - return 0.0, 0.0 - input_cost: Final = float(info.get("input_cost_per_token") or 0.0) - cache_read_cost: Final = info.get("cache_read_input_token_cost") - if cache_read_cost is None: - return input_cost, input_cost - return input_cost, float(cache_read_cost) + 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, + ) class _ModelIdentity(NamedTuple): @@ -434,10 +447,28 @@ def compute_savings_spend( Dollar savings for one request, split by optimization driver. Compression savings price the tokens compression removed at the model's - input rate. Prompt-caching savings price the cache-read tokens at the - difference between the input rate and the discounted cache-read rate; the - read count is derived here from ``usage_object`` so no caller can hand in a - count that disagrees with the usage record. Auto-router savings compare the + input rate. Prompt-caching savings are NET: the cache-read discount minus the + 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. Auto-router savings compare the served ``model`` against the counterfactual baseline the router recorded on its ``routing_decision``, and are zero unless the two differ. That record also says whether the conversation was already underway, which is what tells @@ -454,10 +485,21 @@ def compute_savings_spend( 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. """ - input_cost, cache_read_cost = _input_and_cache_read_cost(model, custom_llm_provider) + # Deployment rates when the request came through one, public rates otherwise -- + # `_effective_model_info` merges a deployment's configured prices over the built-in + # map, so a negotiated price is not silently replaced by the list rate. + router_instance: Router | None = llm_router() if llm_router else None + identity: Final = _resolve_model(model, custom_llm_provider) + 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) compression: Final = max(compression_saved_tokens, 0) * input_cost cache_read_input_tokens: Final = extract_cache_read_tokens(usage_object) - prompt_caching: Final = max(cache_read_input_tokens, 0) * max(input_cost - cache_read_cost, 0.0) + 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) if usage is None or not model: @@ -480,9 +522,7 @@ def compute_savings_spend( # Absent means the router never recorded a shape, which is the conservative # reading: charge the cache write rather than claim a first turn's saving. conversation_continuing=decision.get("conversation_continuing") is not False, - selected_info=_effective_model_info( - (router_instance := llm_router() if llm_router else None), model_id, model or "" - ), + 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, ) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 08bb8698cac..b3feb5bd8d6 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -21,6 +21,7 @@ from litellm.proxy.config_resolvers.sso import ( SSO_SECRET_FIELDS, resolve_sso_config, ) +from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.organization_repository import OrganizationRepository @@ -307,6 +308,27 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = { "enable_chat_ui", } +ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution" + +# UI settings derived from the deployment environment. Deliberately kept out of +# ALLOWED_UI_SETTINGS_FIELDS: they are read-only, never persisted, and PATCH +# rejects them so an admin cannot flip an env-gated feature at runtime. +_DERIVED_UI_SETTINGS_FIELDS: Final[frozenset[str]] = frozenset({ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING}) + + +def _derived_ui_setting_value(key: str) -> object: + """The environment-derived value GET reports for ``key``. + + PATCH compares against this rather than rejecting the key outright, so the body GET + hands back is still a valid PATCH body. Rejecting on presence broke read-modify-write: + a client that edited one setting and sent the rest back unchanged got a 400 and lost + the edit it actually wanted. + """ + if key == ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: + return is_ptu_cost_attribution_enabled() + return None + + # Flags that must be synced from the persisted UISettings into # general_settings at runtime (on both read and write). _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [ @@ -1345,21 +1367,15 @@ async def get_ui_settings(): detail={"error": "Database not connected. Please connect a database."}, ) - ui_settings: Mapping[str, JsonValue] = {} - db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( where={"id": "ui_settings"} ) - if db_record and db_record.ui_settings: - ui_settings_json: Final = db_record.ui_settings - if isinstance(ui_settings_json, str): - ui_settings = json.loads(ui_settings_json) - else: - ui_settings = dict(ui_settings_json) + stored: Final = (db_record.ui_settings if db_record else None) or "{}" + parsed: Final = json.loads(stored) if isinstance(stored, str) else stored # Sanitize any unexpected keys from persisted config before returning - ui_settings = {k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS} + ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS} # Sync runtime flags into general_settings so the proxy picks them up # at runtime (covers server restart scenarios). @@ -1377,11 +1393,18 @@ async def get_ui_settings(): # Build config-like object for schema helper config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}} - return await _get_settings_with_schema( + settings: Final = await _get_settings_with_schema( settings_key="ui_settings", settings_class=_get_effective_ui_settings_class(), config=config, ) + return UISettingsResponse( + values={ + **settings["values"], + ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), + }, + field_schema=settings["field_schema"], + ) @router.patch( @@ -1418,6 +1441,20 @@ async def update_ui_settings( detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) + conflicting_keys: Final = sorted( + key + for key, value in settings_body.items() + if key in _DERIVED_UI_SETTINGS_FIELDS and value != _derived_ui_setting_value(key) + ) + if conflicting_keys: + raise HTTPException( + status_code=400, + detail=( + f"Setting(s) {conflicting_keys} are derived from the deployment environment " + "and cannot be changed from the UI." + ), + ) + # Validate against the same effective class GET advertises, so # enterprise-registered fields are typed consistently on both sides. effective_cls: Final = _get_effective_ui_settings_class() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e59c6adaf22..dd0c57aa911 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -544,6 +544,11 @@ class ProxyLogging: for hook in PROXY_HOOKS: proxy_hook = get_proxy_hook(hook) expected_args = inspect.getfullargspec(proxy_hook).args + if "prisma_client" in expected_args and prisma_client is None: + verbose_proxy_logger.debug( + "Skipping proxy hook %s: it requires a database and no prisma client is configured", hook + ) + continue passed_in_args: dict[str, Any] = {} if "internal_usage_cache" in expected_args: passed_in_args["internal_usage_cache"] = self.internal_usage_cache @@ -3481,13 +3486,15 @@ class PrismaClient: r.expires = r.expires.isoformat() elif query_type == "find_all" and expires is not None and reset_at is not None: response = await VerificationTokenRepository(self).table.find_many( + take=limit, where={ "OR": [ {"expires": None}, {"expires": {"gt": expires}}, ], "budget_reset_at": {"lt": reset_at}, - } + "NOT": {"budget_duration": None}, + }, ) if response is not None and len(response) > 0: for r in response: @@ -3537,6 +3544,7 @@ class PrismaClient: response = await UserRepository(self).table.find_many(where=key_val) elif query_type == "find_all" and reset_at is not None: response = await UserRepository(self).table.find_many( + take=limit, where={ # A user seeded from default_internal_user_params # (or created via /user/new without an explicit @@ -3547,16 +3555,12 @@ class PrismaClient: # of the row, silently exceeding max_budget. Treat a # NULL budget_reset_at with a non-NULL budget_duration # as due, matching the budget-table query below. + "NOT": {"budget_duration": None}, "OR": [ - { - "AND": [ - {"budget_reset_at": None}, - {"NOT": {"budget_duration": None}}, - ] - }, + {"budget_reset_at": None}, {"budget_reset_at": {"lt": reset_at}}, ], - } + }, ) elif query_type == "find_all" and user_id_list is not None: response = await UserRepository(self).table.find_many(where={"user_id": {"in": user_id_list}}) @@ -3612,17 +3616,14 @@ class PrismaClient: elif table_name == "budget" and reset_at is not None: if query_type == "find_all": response = await BudgetRepository(self).table.find_many( + take=limit, where={ + "NOT": {"budget_duration": None}, "OR": [ - { - "AND": [ - {"budget_reset_at": None}, - {"NOT": {"budget_duration": None}}, - ] - }, + {"budget_reset_at": None}, {"budget_reset_at": {"lt": reset_at}}, - ] - } + ], + }, ) return response @@ -3640,20 +3641,17 @@ class PrismaClient: ) elif query_type == "find_all" and reset_at is not None: response = await TeamRepository(self).table.find_many( + take=limit, where={ # Same NULL budget_reset_at gap as the user query # above: a team with a budget_duration but no # initialized budget_reset_at would never be reset. + "NOT": {"budget_duration": None}, "OR": [ - { - "AND": [ - {"budget_reset_at": None}, - {"NOT": {"budget_duration": None}}, - ] - }, + {"budget_reset_at": None}, {"budget_reset_at": {"lt": reset_at}}, ], - } + }, ) elif query_type == "find_all" and user_id is not None: response = await TeamRepository(self).table.find_many( diff --git a/litellm/proxy/vector_store_endpoints/endpoints.py b/litellm/proxy/vector_store_endpoints/endpoints.py index c2483c81d6c..b497247f576 100644 --- a/litellm/proxy/vector_store_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_endpoints/endpoints.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Annotated, Any, Final from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -18,7 +18,8 @@ from litellm.proxy.vector_store_endpoints.utils import ( get_litellm_managed_vector_store, ) from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository -from litellm.types.vector_stores import IndexCreateRequest +from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse +from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry router: Final = APIRouter() ######################################################## @@ -549,14 +550,15 @@ async def index_create( Create an index. Just writes the index to the database. ```bash - curl -L -X POST 'http://0.0.0.0:4000/indexes/create' \ + curl -L -X POST 'http://0.0.0.0:4000/v1/indexes' \ -H 'Content-Type: application/json' \ -H 'Authorization: Bearer sk-1234' \ - -H 'LiteLLM-Beta: indexes_beta=v1' \ - -d '{ + -d '{ "index_name": "dall-e-3", - "vector_store_index": "real-index-name", - "vector_store_name": "azure-ai-search" + "litellm_params": { + "vector_store_index": "real-index-name", + "vector_store_name": "azure-ai-search" + } }' ``` """ @@ -592,3 +594,36 @@ async def index_create( new_index = await ManagedVectorStoreIndexRepository(prisma_client).table.create(data=jsonify_object(index_data)) return new_index.model_dump() + + +@router.get( + "/v1/indexes", + dependencies=[Depends(user_api_key_auth)], + response_model=IndexListResponse, +) +async def index_list( + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> IndexListResponse: + """ + List all vector store indexes. Proxy admin only. + + ```bash + curl -L -X GET 'http://0.0.0.0:4000/v1/indexes' \ + -H 'Authorization: Bearer sk-1234' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + assert_proxy_admin_for_vector_store_index_management( + user_api_key_dict, + operation="list", + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + + indexes: Final = await VectorStoreIndexRegistry._get_vector_store_indexes_from_db(prisma_client) + return IndexListResponse(data=indexes) diff --git a/litellm/proxy/vector_store_endpoints/utils.py b/litellm/proxy/vector_store_endpoints/utils.py index 402fba65558..94ba7c06cad 100644 --- a/litellm/proxy/vector_store_endpoints/utils.py +++ b/litellm/proxy/vector_store_endpoints/utils.py @@ -41,7 +41,7 @@ def _is_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: def assert_proxy_admin_for_vector_store_index_management( user_api_key_dict: UserAPIKeyAuth, *, - operation: Literal["create", "delete", "update"] = "create", + operation: Literal["create", "delete", "update", "list"] = "create", ) -> None: """Raise 403 unless the caller is a proxy admin.""" if _is_proxy_admin(user_api_key_dict): diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 36f1e4cf480..2a9bda08325 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -17,7 +17,8 @@ from __future__ import annotations import hashlib import uuid -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, TypedDict import litellm from litellm._logging import verbose_logger @@ -35,10 +36,32 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion if TYPE_CHECKING: + import httpx + from litellm import Router from litellm.types.rag import RAGIngestOptions +class S3VectorDataPayload(TypedDict): + float32: Sequence[float] + + +class S3VectorEntry(TypedDict): + key: str + data: S3VectorDataPayload + metadata: Mapping[str, str] + + +class S3VectorsQueryMatch(TypedDict, total=False): + key: str + distance: float + metadata: Mapping[str, str] + + +class S3VectorsQueryResponse(TypedDict, total=False): + vectors: Sequence[S3VectorsQueryMatch] + + class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): """ S3 Vectors RAG ingestion using httpx + AWS SigV4 signing. @@ -66,10 +89,10 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): BaseAWSLLM.__init__(self) # Extract config - self.vector_bucket_name = self.vector_store_config["vector_bucket_name"] - self.index_name = self.vector_store_config.get("index_name") - self.distance_metric = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC) - self.non_filterable_metadata_keys = self.vector_store_config.get( + self.vector_bucket_name: str = self.vector_store_config["vector_bucket_name"] + self.index_name: str | None = self.vector_store_config.get("index_name") + self.distance_metric: str = self.vector_store_config.get("distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC) + self.non_filterable_metadata_keys: Sequence[str] = self.vector_store_config.get( "non_filterable_metadata_keys", S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, ) @@ -78,7 +101,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): self.dimension = self._get_dimension_from_config() # Get AWS region using BaseAWSLLM method - _aws_region: Final = self.vector_store_config.get("aws_region_name") + _aws_region: Final[str | None] = self.vector_store_config.get("aws_region_name") self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls( aws_region_name=str(_aws_region) if _aws_region else None ) @@ -135,7 +158,8 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): Returns None if dimension should be auto-detected. """ if "dimension" in self.vector_store_config: - return int(self.vector_store_config["dimension"]) + configured_dimension: Final[int] = self.vector_store_config["dimension"] + return int(configured_dimension) return None async def _ensure_config_initialized(self): @@ -258,7 +282,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): get_body: Final = safe_dumps({"vectorBucketName": self.vector_bucket_name}) try: - response = await self._sign_and_execute_request("POST", get_url, data=get_body) + response: httpx.Response = await self._sign_and_execute_request("POST", get_url, data=get_body) if response.status_code == 200: verbose_logger.debug("Vector bucket %s exists", self.vector_bucket_name) return @@ -294,7 +318,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): get_body: Final = safe_dumps({"vectorBucketName": self.vector_bucket_name, "indexName": self.index_name}) try: - response = await self._sign_and_execute_request("POST", get_url, data=get_body) + response: httpx.Response = await self._sign_and_execute_request("POST", get_url, data=get_body) if response.status_code == 200: verbose_logger.debug("Vector index %s exists", self.index_name) return @@ -311,7 +335,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): ) # Prepare index configuration per AWS API docs - index_config: Final = { + index_config: Final[dict[str, object]] = { "vectorBucketName": self.vector_bucket_name, "indexName": self.index_name, "dataType": "float32", @@ -336,7 +360,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): verbose_logger.exception("Error creating vector index: %s", e) raise - async def _put_vectors(self, vectors: list[dict[str, Any]]): + async def _put_vectors(self, vectors: Sequence[S3VectorEntry]): """ Call PutVectors API to store vectors in S3 Vectors. @@ -355,7 +379,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): } try: - response: Final = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body)) + response: Final[httpx.Response] = await self._sign_and_execute_request( + "POST", url, data=safe_dumps(request_body) + ) if response.status_code in (200, 201): verbose_logger.info("Successfully stored %s vectors in index %s", len(vectors), self.index_name) @@ -442,24 +468,18 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): raise ValueError(error_msg) # Prepare vectors for PutVectors API - vectors: Final = [] - for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): - # Build metadata dict - metadata: dict[str, str] = { - "source_text": chunk, # Non-filterable (for reference) - "chunk_index": str(i), # Filterable - } - - if filename: - metadata["filename"] = filename # Filterable - - vector_obj = { - "key": f"{filename}_{i}" if filename else f"chunk_{i}", - "data": {"float32": embedding}, - "metadata": metadata, - } - - vectors.append(vector_obj) + vectors: Final = [ + S3VectorEntry( + key=f"{filename}_{i}" if filename else f"chunk_{i}", + data=S3VectorDataPayload(float32=embedding), + metadata=( + {"source_text": chunk, "chunk_index": str(i), "filename": filename} + if filename + else {"source_text": chunk, "chunk_index": str(i)} + ), + ) + for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)) + ] # Call PutVectors API await self._put_vectors(vectors) @@ -468,7 +488,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): vector_store_id: Final = f"{self.vector_bucket_name}:{self.index_name}" return vector_store_id, filename - async def query_vector_store(self, vector_store_id: str, query: str, top_k: int = 5) -> dict[str, Any] | None: + async def query_vector_store( + self, vector_store_id: str, query: str, top_k: int = 5 + ) -> S3VectorsQueryResponse | None: """ Query S3 Vectors using QueryVectors API. @@ -489,7 +511,7 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): embedding_model: Final = self.embedding_config.get("model", "text-embedding-3-small") response = await litellm.aembedding(model=embedding_model, input=[query]) - query_embedding: Final = response.data[0]["embedding"] + query_embedding: Final[Sequence[float]] = response.data[0]["embedding"] # Call QueryVectors API url: Final = f"https://s3vectors.{self.aws_region_name}.api.aws/QueryVectors" @@ -504,15 +526,18 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): } try: - response = await self._sign_and_execute_request("POST", url, data=safe_dumps(request_body)) + query_response: Final[httpx.Response] = await self._sign_and_execute_request( + "POST", url, data=safe_dumps(request_body) + ) - if response.status_code == 200: - results: Final = response.json() + if query_response.status_code == 200: + results: Final[S3VectorsQueryResponse] = query_response.json() + matches: Final = results.get("vectors") verbose_logger.debug("Query returned %s results", len(results.get("vectors", []))) # Check if query terms appear in results - if results.get("vectors"): - for result in results["vectors"]: + if matches: + for result in matches: metadata = result.get("metadata", {}) source_text = metadata.get("source_text", "") if query.lower() in source_text.lower(): @@ -521,7 +546,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Return results even if exact match not found return results else: - verbose_logger.error("QueryVectors failed with status %s: %s", response.status_code, response.text) + verbose_logger.error( + "QueryVectors failed with status %s: %s", query_response.status_code, query_response.text + ) return None except Exception as e: verbose_logger.exception("Error querying vectors: %s", e) diff --git a/litellm/repositories/__init__.py b/litellm/repositories/__init__.py index 4f020480f9e..e2e7f1fac73 100644 --- a/litellm/repositories/__init__.py +++ b/litellm/repositories/__init__.py @@ -70,10 +70,14 @@ from litellm.repositories.table_repositories import ( ) from litellm.repositories.team_repository import TeamRepository from litellm.repositories.unit_of_work import ( + BudgetCascadeUnitOfWork, + BudgetWindowWrites, KeySpendResetWrites, + LinkedSpendResetWrites, SpendResetUnitOfWork, TeamSpendResetWrites, UserSpendResetWrites, + budget_cascade_unit_of_work, spend_reset_unit_of_work, ) from litellm.repositories.user_repository import UserRepository @@ -88,7 +92,9 @@ __all__ = [ "AgentsRepository", "AuditLogRepository", "BatchTable", + "BudgetCascadeUnitOfWork", "BudgetRepository", + "BudgetWindowWrites", "CacheConfigRepository", "ClaudeCodePluginRepository", "ConfigOverridesRepository", @@ -107,6 +113,7 @@ __all__ = [ "InvitationLinkRepository", "JWTKeyMappingRepository", "KeySpendResetWrites", + "LinkedSpendResetWrites", "MCPServerRepository", "MCPToolsetRepository", "MCPUserCredentialsRepository", @@ -149,5 +156,6 @@ __all__ = [ "WorkflowEventRepository", "WorkflowMessageRepository", "WorkflowRunRepository", + "budget_cascade_unit_of_work", "spend_reset_unit_of_work", ] diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 6aff196ff10..055c68163f9 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -29,6 +29,8 @@ class SpendLinkedTable(Protocol[RowT_co]): class BatchTable(Protocol): def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... + def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... + class PrismaBatch(Protocol): @property @@ -40,4 +42,19 @@ class PrismaBatch(Protocol): @property def litellm_teamtable(self) -> BatchTable: ... + @property + def litellm_budgettable(self) -> BatchTable: ... + + @property + def litellm_teammembership(self) -> BatchTable: ... + + @property + def litellm_organizationtable(self) -> BatchTable: ... + + @property + def litellm_tagtable(self) -> BatchTable: ... + + @property + def litellm_endusertable(self) -> BatchTable: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 682e69d11eb..e504baceb9f 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -1,17 +1,21 @@ """ -Unit of work over a single Prisma batch. +Units of work over a single Prisma batch. -``spend_reset_unit_of_work`` opens one ``db.batch_()`` and binds a typed write +Each context manager here opens one ``db.batch_()`` and binds a typed write repository per table to it, so every update queued through the yielded object lands in the same transaction. The batch commits when the block exits cleanly and is abandoned, writing nothing, when the block raises. -Each write repository queues narrow ``{spend, budget_reset_at}`` updates +``spend_reset_unit_of_work`` covers the per-row key/user/team resets; +``budget_cascade_unit_of_work`` covers a budget tier's reset, where the +dependent spend and the tier's next window have to move together. + +Each write repository queues narrow ``{spend}`` / ``{budget_reset_at}`` updates instead of full-model writes, which trip ``prisma.errors.DataError`` on rows carrying fields the update input type rejects (see #27730). """ -from collections.abc import AsyncGenerator, Callable +from collections.abc import AsyncGenerator, Callable, Mapping from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime @@ -43,6 +47,24 @@ class TeamSpendResetWrites: self.table.update(where={"team_id": team_id}, data={"spend": 0, "budget_reset_at": budget_reset_at}) +@dataclass(frozen=True, slots=True) +class LinkedSpendResetWrites: + table: BatchTable + + def queue_spend_zero(self, where: Mapping[str, object]) -> None: + self.table.update_many(where=where, data={"spend": 0}) + + +@dataclass(frozen=True, slots=True) +class BudgetWindowWrites: + table: BatchTable + + def queue_window_advance(self, budget_id: str, budget_reset_at: datetime) -> None: + """``update_many`` so a tier deleted between the read and the commit is a + no-op row count instead of a P2025 that aborts the whole chunk.""" + self.table.update_many(where={"budget_id": budget_id}, data={"budget_reset_at": budget_reset_at}) + + @dataclass(frozen=True, slots=True) class SpendResetUnitOfWork: keys: KeySpendResetWrites @@ -50,6 +72,23 @@ class SpendResetUnitOfWork: teams: TeamSpendResetWrites +@dataclass(frozen=True, slots=True) +class BudgetCascadeUnitOfWork: + """Every write a budget-tier reset performs, bound to one batch. + + The dependent spend rows and the budget rows' ``budget_reset_at`` advance + must land together: advancing the window without zeroing the spend it + gates leaves the dependents pinned at their cap until the next window. + """ + + team_memberships: LinkedSpendResetWrites + keys: LinkedSpendResetWrites + organizations: LinkedSpendResetWrites + tags: LinkedSpendResetWrites + endusers: LinkedSpendResetWrites + budgets: BudgetWindowWrites + + @asynccontextmanager async def spend_reset_unit_of_work(new_batch: Callable[[], PrismaBatch]) -> AsyncGenerator[SpendResetUnitOfWork, None]: batch = new_batch() @@ -59,3 +98,19 @@ async def spend_reset_unit_of_work(new_batch: Callable[[], PrismaBatch]) -> Asyn teams=TeamSpendResetWrites(table=batch.litellm_teamtable), ) await batch.commit() + + +@asynccontextmanager +async def budget_cascade_unit_of_work( + new_batch: Callable[[], PrismaBatch], +) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]: + batch = new_batch() + yield BudgetCascadeUnitOfWork( + team_memberships=LinkedSpendResetWrites(table=batch.litellm_teammembership), + keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken), + organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable), + tags=LinkedSpendResetWrites(table=batch.litellm_tagtable), + endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable), + budgets=BudgetWindowWrites(table=batch.litellm_budgettable), + ) + await batch.commit() diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 8448db11904..c6e17502e5d 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -2,7 +2,10 @@ import re import traceback from collections.abc import Iterable, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypedDict, overload + +from openai.types.chat import ChatCompletionToolParam +from openai.types.responses.function_tool_param import FunctionToolParam from litellm._logging import verbose_logger from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG @@ -18,6 +21,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamingResponse, ) +from litellm.types.llms.openai import ToolParam as ResponsesToolParam from litellm.types.utils import ( CallTypes, Choices, @@ -36,10 +40,14 @@ else: MCPTool = Any # NOTE: We intentionally keep ToolParam as a broad type here to avoid tight coupling -# to optional OpenAI SDK typing symbols in environments that may not have them available. -# `Any` is used to keep mypy compatible with the broader OpenAI tool union types -# passed around in Responses API while still allowing dict-style access at runtime. -ToolParam = Any +ToolParam: TypeAlias = Mapping[str, object] + + +class MCPToolResult(TypedDict): + tool_call_id: str | None + result: str + name: str | None + LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" @@ -199,13 +207,12 @@ class LiteLLM_Proxy_MCP_Handler: _get_tools_from_mcp_servers, ) - mcp_servers: Final[list[str]] = [] - if mcp_tools_with_litellm_proxy: - for _tool in mcp_tools_with_litellm_proxy: - # if user specifies servers as server_url: litellm_proxy/mcp/zapier,github then return zapier,github - server_url = _tool.get("server_url", "") if isinstance(_tool, dict) else "" - if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX): - mcp_servers.append(server_url.split("/")[-1]) + mcp_servers: Final = [ + server_url.split("/")[-1] + for _tool in (mcp_tools_with_litellm_proxy or ()) + for server_url in (_tool.get("server_url", "") if isinstance(_tool, dict) else "",) + if isinstance(server_url, str) and server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL_PREFIX) + ] # Resolve toolset names: collect all toolset IDs first, then apply their # combined permissions in a single pass so multiple toolsets are unioned @@ -279,15 +286,15 @@ class LiteLLM_Proxy_MCP_Handler: allowed_mcp_servers=allowed_mcp_servers, ) - server_names: Final[list[str]] = [] - for server in allowed_mcp_servers: - if server is None: - continue - server_name = ( - getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None) + server_names: Final = [ + server_name + for server in allowed_mcp_servers + if server is not None + for server_name in ( + getattr(server, "server_name", None) or getattr(server, "alias", None) or getattr(server, "name", None), ) - if isinstance(server_name, str): - server_names.append(server_name) + if isinstance(server_name, str) + ] return tools, server_names @@ -305,8 +312,8 @@ class LiteLLM_Proxy_MCP_Handler: List of deduplicated MCP tools The returned dictionary maps each tool_name to the server_name """ - seen_names: Final = set() - deduplicated_tools: Final = [] + seen_names: Final[set[str]] = set() + deduplicated_tools: Final[list[MCPTool]] = [] tool_server_map: Final[dict[str, str]] = {} for tool in mcp_tools: @@ -331,7 +338,7 @@ class LiteLLM_Proxy_MCP_Handler: ) -> list[MCPTool]: """Filter MCP tools based on allowed_tools parameter from the original tool configs.""" # Collect all allowed tool names from all MCP tool configs - allowed_tool_names: Final = set() + allowed_tool_names: Final[set[str]] = set() for tool_config in mcp_tools_with_litellm_proxy: if isinstance(tool_config, dict) and "allowed_tools" in tool_config: allowed_tools = tool_config.get("allowed_tools", []) @@ -343,23 +350,13 @@ class LiteLLM_Proxy_MCP_Handler: return mcp_tools # Filter tools based on allowed names - filtered_tools: Final = [] - for mcp_tool in mcp_tools: - if isinstance(mcp_tool, dict): - tool_name = mcp_tool.get("name") - else: - tool_name = getattr(mcp_tool, "name", None) - - if not tool_name: - continue - - if tool_name in allowed_tool_names: - filtered_tools.append(mcp_tool) - continue - - unprefixed_name, _ = split_server_prefix_from_name(tool_name) - if unprefixed_name in allowed_tool_names: - filtered_tools.append(mcp_tool) + filtered_tools: Final = [ + mcp_tool + for mcp_tool in mcp_tools + for tool_name in (mcp_tool.get("name") if isinstance(mcp_tool, dict) else getattr(mcp_tool, "name", None),) + if tool_name + and (tool_name in allowed_tool_names or split_server_prefix_from_name(tool_name)[0] in allowed_tool_names) + ] return filtered_tools @@ -448,24 +445,37 @@ class LiteLLM_Proxy_MCP_Handler: return deduplicated_mcp_tools, tool_server_map + @overload + @staticmethod + def _transform_mcp_tools_to_openai( + mcp_tools: Sequence[MCPTool], + target_format: Literal["responses"] = ..., + ) -> list[FunctionToolParam]: ... + + @overload + @staticmethod + def _transform_mcp_tools_to_openai( + mcp_tools: Sequence[MCPTool], + target_format: Literal["chat"], + ) -> list[ChatCompletionToolParam]: ... + @staticmethod def _transform_mcp_tools_to_openai( mcp_tools: Sequence[MCPTool], target_format: Literal["responses", "chat"] = "responses", - ) -> list[Any]: + ) -> Sequence[FunctionToolParam | ChatCompletionToolParam]: """Transform MCP tools to OpenAI-compatible format.""" from litellm.experimental_mcp_client.tools import ( transform_mcp_tool_to_openai_responses_api_tool, transform_mcp_tool_to_openai_tool, ) - openai_tools: Final[list[Any]] = [] - for mcp_tool in mcp_tools: - if target_format == "chat": - openai_tool = transform_mcp_tool_to_openai_tool(mcp_tool) - else: - openai_tool = transform_mcp_tool_to_openai_responses_api_tool(mcp_tool) - openai_tools.append(openai_tool) + openai_tools: Final = [ + transform_mcp_tool_to_openai_tool(mcp_tool) + if target_format == "chat" + else transform_mcp_tool_to_openai_responses_api_tool(mcp_tool) + for mcp_tool in mcp_tools + ] return openai_tools @@ -496,9 +506,9 @@ class LiteLLM_Proxy_MCP_Handler: return True @staticmethod - def _extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[Any]: + def _extract_tool_calls_from_response(response: ResponsesAPIResponse) -> list[object]: """Extract tool calls from the response output.""" - tool_calls: Final[list[Any]] = [] + tool_calls: Final[list[object]] = [] for output_item in response.output: # Check if this is a function call output item if isinstance(output_item, dict) and output_item.get("type") == "function_call": @@ -533,7 +543,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _extract_tool_call_details( - tool_call, + tool_call: object, ) -> tuple[str | None, str | None, str | None]: """Extract tool name, arguments, and call_id from a tool call.""" if isinstance(tool_call, dict): @@ -566,7 +576,7 @@ class LiteLLM_Proxy_MCP_Handler: return tool_name, tool_arguments, tool_call_id @staticmethod - def _parse_tool_arguments(tool_arguments: Any) -> dict[str, Any]: + def _parse_tool_arguments(tool_arguments: str | None) -> dict[str, object]: """Parse tool arguments, handling both string and dict formats.""" import json @@ -591,23 +601,18 @@ class LiteLLM_Proxy_MCP_Handler: # Fallback to generic handling if MCP types not available return "Tool executed successfully" - text_parts: Final = [] - other_content_types: Final = [] - - for content_item in result.content: - if isinstance(content_item, TextContent): - # Text content - extract the text - text_parts.append(str(content_item.text)) - elif isinstance(content_item, ImageContent): - # Image content - other_content_types.append("Image") - elif isinstance(content_item, EmbeddedResource): - # Embedded resource - other_content_types.append("EmbeddedResource") - else: - # Other unknown content types - content_type = type(content_item).__name__ - other_content_types.append(content_type) + text_parts: Final = [ + str(content_item.text) for content_item in result.content if isinstance(content_item, TextContent) + ] + other_content_types: Final = [ + "Image" + if isinstance(content_item, ImageContent) + else "EmbeddedResource" + if isinstance(content_item, EmbeddedResource) + else type(content_item).__name__ + for content_item in result.content + if not isinstance(content_item, TextContent) + ] # Combine text parts if any result_text = " ".join(text_parts) if text_parts else "" @@ -631,7 +636,7 @@ class LiteLLM_Proxy_MCP_Handler: litellm_call_id: str | None = None, litellm_trace_id: str | None = None, request_tags: list[str] | None = None, - ) -> list[dict[str, Any]]: + ) -> list[MCPToolResult]: """Execute tool calls and return results.""" from fastapi import HTTPException @@ -645,11 +650,11 @@ class LiteLLM_Proxy_MCP_Handler: ) from litellm.proxy.proxy_server import proxy_logging_obj - tool_results: Final = [] + tool_results: Final[list[MCPToolResult]] = [] tool_call_id: str | None = None rules_obj: Final = Rules() for tool_call in tool_calls: - logging_request_data: dict[str, Any] = {} + logging_request_data: dict[str, object] = {} tool_name: str | None = None try: ( @@ -678,7 +683,7 @@ class LiteLLM_Proxy_MCP_Handler: sanitized_tool_name = strip_known_server_prefix(resolved_tool_name, mcp_server) start_time = datetime.now() - logging_input = [ + logging_input: Sequence[Mapping[str, object]] = [ { "role": "tool", "content": { @@ -688,13 +693,14 @@ class LiteLLM_Proxy_MCP_Handler: } ] tool_logging_call_id = litellm_call_id or str(uuid.uuid4()) + logging_metadata: dict[str, object] = { + "tool_call_id": tool_call_id, + "tool_name": sanitized_tool_name, + "server_name": server_name, + } logging_request_data = { "model": f"MCP: {tool_name}", - "metadata": { - "tool_call_id": tool_call_id, - "tool_name": sanitized_tool_name, - "server_name": server_name, - }, + "metadata": logging_metadata, "input": logging_input, "call_type": CallTypes.call_mcp_tool.value, "litellm_call_id": tool_logging_call_id, @@ -712,7 +718,7 @@ class LiteLLM_Proxy_MCP_Handler: if litellm_trace_id: logging_request_data["litellm_trace_id"] = litellm_trace_id if request_tags: - logging_request_data["metadata"]["tags"] = request_tags + logging_metadata["tags"] = request_tags if user_api_key_auth is not None: from litellm.proxy.litellm_pre_call_utils import ( LiteLLMProxyRequestSetup, @@ -902,16 +908,16 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _create_follow_up_messages_for_chat( - original_messages: list[Any], + original_messages: list[object], response: ModelResponse, tool_results: Sequence[Mapping[str, object]], - ) -> list[Any]: + ) -> Sequence[Mapping[str, object]]: """Create follow-up chat messages that include tool execution results.""" from copy import deepcopy from litellm.utils import convert_list_message_to_dict - follow_up_messages: list[Any] = convert_list_message_to_dict(deepcopy(original_messages)) + follow_up_messages: list[dict[str, object]] = convert_list_message_to_dict(deepcopy(original_messages)) if not follow_up_messages: follow_up_messages = [] @@ -950,9 +956,9 @@ class LiteLLM_Proxy_MCP_Handler: response: ResponsesAPIResponse, tool_results: Sequence[Mapping[str, object]], original_input: str | ResponseInputParam | None = None, - ) -> list[Any]: + ) -> list[object]: """Create follow-up input with tool results in proper format.""" - follow_up_input: Final[list[Any]] = [] + follow_up_input: Final[list[object]] = [] # Add original user input if available to maintain conversation context if original_input: @@ -964,8 +970,8 @@ class LiteLLM_Proxy_MCP_Handler: follow_up_input.append(original_input) # Add the assistant message with function calls - assistant_message_content: Final[list[Any]] = [] - function_calls: Final[list[dict[str, Any]]] = [] + assistant_message_content: Final[list[object]] = [] + function_calls: Final[list[dict[str, object]]] = [] for output_item in response.output: if not isinstance(output_item, dict) and hasattr(output_item, "model_dump"): @@ -1027,7 +1033,7 @@ class LiteLLM_Proxy_MCP_Handler: async def _make_follow_up_call( follow_up_input: list[Any], model: str, - all_tools: list[Any] | None, + all_tools: Sequence[ResponsesToolParam] | None, response_id: str, **call_params: Any, ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: @@ -1044,7 +1050,7 @@ class LiteLLM_Proxy_MCP_Handler: async def _log_mcp_tool_failure( *, proxy_logging_obj: Optional["ProxyLogging"], - user_api_key_auth: Any, + user_api_key_auth: "UserAPIKeyAuth | None", request_data: dict[str, object], error: Exception, ) -> None: @@ -1072,7 +1078,7 @@ class LiteLLM_Proxy_MCP_Handler: all_tools: Sequence[object] | None, mcp_tools_with_litellm_proxy: list[Mapping[str, object]], mcp_discovery_events: list[ResponsesAPIStreamingResponse], - call_params: dict[str, Any], + call_params: Mapping[str, object], previous_response_id: str | None, tool_server_map: dict[str, str], **kwargs, @@ -1115,10 +1121,10 @@ class LiteLLM_Proxy_MCP_Handler: input: str | ResponseInputParam, model: str, all_tools: Sequence[object] | None, - call_params: dict[str, Any], + call_params: Mapping[str, object], previous_response_id: str | None, - **kwargs, - ) -> dict[str, Any]: + **kwargs: object, + ) -> dict[str, object]: """ Build a clean request parameters dictionary for MCP streaming. @@ -1126,7 +1132,7 @@ class LiteLLM_Proxy_MCP_Handler: in a clean, maintainable way. """ # Start with the core required parameters - request_params: Final = { + request_params: Final[dict[str, object]] = { "input": input, "model": model, "tools": all_tools, @@ -1146,7 +1152,7 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def _create_tool_execution_events( - tool_calls: Sequence[object], tool_results: list[dict[str, Any]] + tool_calls: Sequence[object], tool_results: Sequence[MCPToolResult] ) -> list[ResponsesAPIStreamingResponse]: """ Create MCP tool execution events for streaming. diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index e4cc36de06c..186852f91c2 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -19,13 +19,13 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, - ToolParam, ) if TYPE_CHECKING: from mcp.types import Tool as MCPTool from litellm.proxy._types import UserAPIKeyAuth + from litellm.responses.mcp.litellm_proxy_mcp_handler import MCPToolResult else: MCPTool = Any @@ -33,7 +33,7 @@ MAX_MCP_TOOL_CALL_ROUNDS: Final = 5 async def create_mcp_list_tools_events( - mcp_tools_with_litellm_proxy: list[ToolParam], + mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]], user_api_key_auth: "UserAPIKeyAuth | None", base_item_id: str, pre_processed_mcp_tools: list[MCPTool], @@ -44,13 +44,14 @@ async def create_mcp_list_tools_events( try: # Extract MCP server names - mcp_servers: Final = [] - for tool in mcp_tools_with_litellm_proxy: - if isinstance(tool, dict) and "server_url" in tool: - server_url = tool.get("server_url") - if isinstance(server_url, str) and server_url.startswith("litellm_proxy/mcp/"): - server_name = server_url.split("/")[-1] - mcp_servers.append(server_name) + _mcp_servers: Final = [ + server_url.split("/")[-1] + for tool in mcp_tools_with_litellm_proxy + if isinstance(tool, dict) + and "server_url" in tool + and isinstance(server_url := tool.get("server_url"), str) + and server_url.startswith("litellm_proxy/mcp/") + ] # Emit list tools in progress event in_progress_event: Final = MCPListToolsInProgressEvent( @@ -65,15 +66,14 @@ async def create_mcp_list_tools_events( filtered_mcp_tools: Final = pre_processed_mcp_tools # Convert tools to dict format for the event - mcp_tools_dict: Final = [] - for tool in filtered_mcp_tools: - if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")): - # Type cast to help mypy understand this is safe after hasattr check - mcp_tools_dict.append(cast(Any, tool).model_dump()) - elif hasattr(tool, "__dict__"): - mcp_tools_dict.append(tool.__dict__) - else: - mcp_tools_dict.append({"name": getattr(tool, "name", str(tool))}) + _mcp_tools_dict: Final = [ + tool.model_dump() + if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")) + else tool.__dict__ + if hasattr(tool, "__dict__") + else {"name": getattr(tool, "name", str(tool))} + for tool in filtered_mcp_tools + ] # Emit list tools completed event completed_event: Final = MCPListToolsCompletedEvent( @@ -96,21 +96,18 @@ async def create_mcp_list_tools_events( server_label = str(server_label_value) if server_label_value is not None else "" # Format tools for OpenAI output_item.done format - formatted_tools: Final = [] - for tool in filtered_mcp_tools: - tool_dict = { + formatted_tools: Final = [ + { "name": getattr(tool, "name", "unknown"), "description": getattr(tool, "description", ""), "annotations": {"read_only": False}, + **dict.fromkeys( + ("input_schema",) if hasattr(tool, "inputSchema") or hasattr(tool, "input_schema") else (), + getattr(tool, "inputSchema", getattr(tool, "input_schema", None)), + ), } - - # Add input_schema if available - if hasattr(tool, "inputSchema"): - tool_dict["input_schema"] = getattr(tool, "inputSchema") - elif hasattr(tool, "input_schema"): - tool_dict["input_schema"] = getattr(tool, "input_schema") - - formatted_tools.append(tool_dict) + for tool in filtered_mcp_tools + ] # Create the output_item.done event with MCP tools list output_item_done_event = OutputItemDoneEvent( @@ -166,7 +163,7 @@ async def create_mcp_list_tools_events( def create_mcp_call_events( tool_name: str, - tool_call_id: str, + tool_call_id: str | None, arguments: str, result: str | None = None, base_item_id: str | None = None, @@ -256,9 +253,12 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): 4. Emits tool execution events in the stream """ + model: str + tool_results: "Sequence[MCPToolResult]" + def __init__( self, - base_iterator: Any, # Can be None - will be created internally + base_iterator: "BaseResponsesAPIStreamingIterator | ResponsesAPIResponse | None", # created internally when None mcp_events: list[ResponsesAPIStreamingResponse], tool_server_map: dict[str, str], mcp_tools_with_litellm_proxy: Sequence[Mapping[str, object]] | None = None, @@ -285,7 +285,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.tool_server_map = tool_server_map # Iterator references - self.base_iterator: Any | ResponsesAPIResponse | None = base_iterator # Will be created when needed + self.base_iterator: BaseResponsesAPIStreamingIterator | ResponsesAPIResponse | None = ( + base_iterator # Will be created when needed + ) # Response collection for tool execution self.collected_response: ResponsesAPIResponse | None = None diff --git a/litellm/router.py b/litellm/router.py index feaf69a44ae..9cece2014af 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -22,6 +22,7 @@ import weakref from collections import defaultdict from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence from functools import lru_cache +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeVar, Union, cast import anyio @@ -111,12 +112,14 @@ from litellm.router_utils.common_utils import ( filter_web_search_deployments, resolve_model_group_alias, truncate_fallback_error_detail, + warn_on_provider_credential_mismatch, ) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( DEFAULT_COOLDOWN_TIME_SECONDS, _async_get_cooldown_deployments, _async_get_cooldown_deployments_with_debug_info, + _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper across router_utils submodules, matching the other cooldown_handlers imports on this line _get_cooldown_deployments, _set_cooldown_deployments, is_advisor_orchestration_failure, @@ -1898,7 +1901,7 @@ class Router: # Set per-deployment num_retries on exception for retry logic if deployment is not None: self._set_deployment_num_retries_on_exception(e, deployment) - self._set_failed_deployment_id_on_exception(e, deployment) + self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) raise e def _get_silent_experiment_kwargs(self, **kwargs) -> dict: @@ -2961,7 +2964,7 @@ class Router: # Set per-deployment num_retries on exception for retry logic if deployment is not None: self._set_deployment_num_retries_on_exception(e, deployment) - self._set_failed_deployment_id_on_exception(e, deployment) + self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) raise e except Exception as e: verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) @@ -2970,7 +2973,7 @@ class Router: # Set per-deployment num_retries on exception for retry logic if deployment is not None: self._set_deployment_num_retries_on_exception(e, deployment) - self._set_failed_deployment_id_on_exception(e, deployment) + self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) raise e def _update_kwargs_before_fallbacks( @@ -3016,7 +3019,7 @@ class Router: except (ValueError, TypeError): pass # Skip if value can't be converted to int - def _set_failed_deployment_id_on_exception(self, exception: Exception, deployment: dict) -> None: + def _set_failed_deployment_id_on_exception(self, exception: Exception, deployment: Mapping[str, Any]) -> None: """ Stamp the failed deployment's `model_info.id` on the exception so the fallback layer can exclude it from subsequent re-picks within the same @@ -3035,6 +3038,16 @@ class Router: except Exception: pass + def _stamp_failed_deployment_id_with_effective_model_info( + self, exception: Exception, deployment: Mapping[str, Any], kwargs: Mapping[str, Any] + ) -> None: + # A client-side-credential call gets a dynamic deployment id generated inside + # _update_kwargs_with_deployment and stamped into kwargs["model_info"]; stamping + # the static shared deployment's id instead would let one tenant's bad credentials + # cool down the deployment every other tenant sharing this config relies on. + effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({}) + self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info})) + def _update_kwargs_with_default_litellm_params( self, kwargs: dict, metadata_variable_name: str | None = "metadata" ) -> None: @@ -4521,10 +4534,11 @@ class Router: passthrough_on_no_deployment: Final = kwargs.pop("passthrough_on_no_deployment", False) function_name: Final = "_ageneric_api_call_with_fallbacks" + deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) try: - deployment: Final = await self.async_get_available_deployment( + deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init above model=model, request_kwargs=kwargs, messages=kwargs.get("messages", None), @@ -4601,6 +4615,8 @@ class Router: ) if model is not None: self.fail_calls[model] += 1 + if deployment is not None: + self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) raise e async def _aresponses_with_streaming_fallbacks( @@ -7078,7 +7094,9 @@ class Router: ) # Determine cooldown time with priority: deployment config > response header > router default - deployment_cooldown: Final = litellm_params.get("cooldown_time", None) + deployment_cooldown: Final = _first_present( + _model_info if isinstance(_model_info, dict) else None, litellm_params, key="cooldown_time" + ) header_cooldown = None if exception_headers is not None: @@ -7523,6 +7541,7 @@ class Router: """ try: litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(**_litellm_params) + warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=_litellm_params) deployment = Deployment( **deployment_info, model_name=_model_name, @@ -8215,6 +8234,11 @@ class Router: if _deployment_model_id and self.has_model_id(_deployment_model_id): return None + warn_on_provider_credential_mismatch( + model_name=deployment.model_name, + litellm_params=deployment.litellm_params.model_dump(exclude_none=True), + ) + # add to model list _deployment: Final = deployment.to_json(exclude_none=True) # initialize client @@ -11800,6 +11824,23 @@ class Router: and allowed_fails_policy.BadRequestErrorAllowedFails is not None ): return allowed_fails_policy.BadRequestErrorAllowedFails + if ( + isinstance(exception, litellm.InternalServerError) + and allowed_fails_policy.InternalServerErrorAllowedFails is not None + ): + return allowed_fails_policy.InternalServerErrorAllowedFails + if ( + isinstance(exception, litellm.ServiceUnavailableError) + and allowed_fails_policy.ServiceUnavailableErrorAllowedFails is not None + ): + return allowed_fails_policy.ServiceUnavailableErrorAllowedFails + if ( + isinstance(exception, litellm.BadGatewayError) + and allowed_fails_policy.BadGatewayErrorAllowedFails is not None + ): + return allowed_fails_policy.BadGatewayErrorAllowedFails + if isinstance(exception, litellm.NotFoundError) and allowed_fails_policy.NotFoundErrorAllowedFails is not None: + return allowed_fails_policy.NotFoundErrorAllowedFails def _initialize_alerting(self): from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 6fad2dd31e9..280a7defcf8 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -1,15 +1,18 @@ import hashlib import json from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject -from litellm._logging import verbose_logger +from litellm._logging import verbose_logger, verbose_router_logger from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.types.router import CredentialLiteLLMParams +from litellm.types.utils import LlmProviders def _is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool: @@ -210,3 +213,77 @@ def filter_web_search_deployments( if len(healthy_deployments) > 0 and len(final_deployments) == 0: verbose_logger.warning("No deployments support web search for request") return final_deployments + + +# Credential params that only one provider family reads, paired with the providers +# that read them. A deployment carrying them while resolving elsewhere is almost +# always a missing route prefix: `model: claude-sonnet-5` with `aws_region_name` +# set resolves to the first-party Anthropic API, silently ignores the AWS +# credentials, and 401s at request time. +_AWS_PROVIDERS: Final = frozenset( + provider.value for provider in LlmProviders if provider.value.startswith(("bedrock", "sagemaker")) +) +_VERTEX_PROVIDERS: Final = frozenset( + provider.value for provider in LlmProviders if provider.value.startswith("vertex_ai") +) + +PROVIDER_SCOPED_CREDENTIAL_PARAMS: Final[Mapping[str, frozenset[str]]] = MappingProxyType( + { + "aws_access_key_id": _AWS_PROVIDERS, + "aws_profile_name": _AWS_PROVIDERS, + "aws_region_name": _AWS_PROVIDERS, + "aws_role_name": _AWS_PROVIDERS, + "aws_secret_access_key": _AWS_PROVIDERS, + "aws_session_name": _AWS_PROVIDERS, + "aws_session_token": _AWS_PROVIDERS, + "aws_web_identity_token": _AWS_PROVIDERS, + "vertex_credentials": _VERTEX_PROVIDERS, + "vertex_location": _VERTEX_PROVIDERS, + "vertex_project": _VERTEX_PROVIDERS, + } +) + + +def warn_on_provider_credential_mismatch(model_name: str, litellm_params: Mapping[str, object]) -> str | None: + """ + Warn when a deployment carries one provider's credentials but resolves to another. + + Returns the warning text (for tests), or None when the deployment is consistent + or its provider cannot be resolved. Never raises: a deployment litellm cannot + classify is left alone rather than blocking router startup. + + Only inline credential params are examined. A deployment that sources them + through ``litellm_credential_name`` resolves them after registration, so it + carries none of these keys here and is left alone rather than warned about + on incomplete information. + """ + model: Final = litellm_params.get("model") + if not isinstance(model, str) or not model: + return None + scoped: Final = tuple(param for param in PROVIDER_SCOPED_CREDENTIAL_PARAMS if litellm_params.get(param) is not None) + if not scoped: + return None + custom_llm_provider: Final = litellm_params.get("custom_llm_provider") + try: + _, resolved_provider, _, _ = get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None, + ) + except BadRequestError: + return None + mismatched: Final = sorted( + param for param in scoped if resolved_provider not in PROVIDER_SCOPED_CREDENTIAL_PARAMS[param] + ) + if not mismatched: + return None + expected: Final = sorted( + {provider for param in mismatched for provider in PROVIDER_SCOPED_CREDENTIAL_PARAMS[param]} + ) + warning: Final = ( + f"Deployment '{model_name}' sets {mismatched} but 'model={model}' resolves to provider " + f"'{resolved_provider}', which ignores them. Those params are read by {expected}, so this is " + f"usually a missing route prefix (e.g. '{expected[0]}/{model}'); as written the request goes to " + f"'{resolved_provider}' and will fail on that provider's credentials." + ) + verbose_router_logger.warning(warning) + return warning diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index 8d3b897ae3e..9e7f457f631 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -4,6 +4,7 @@ Wrapper around router cache. Meant to handle model cooldown logic import functools import time +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from typing_extensions import TypedDict @@ -28,6 +29,12 @@ class CooldownCacheValue(TypedDict): cooldown_time: float +# Cap on the corrected in-memory TTL set in `_corrected_active_cooldown`: re-checks the +# real remaining cooldown against Redis at least this often, so an entry that later gets +# deleted or extended in Redis before its original deadline is still noticed promptly. +_MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0 + + class CooldownCache: def __init__(self, cache: DualCache, default_cooldown_time: float): self.cache = cache @@ -100,6 +107,30 @@ class CooldownCache: def get_cooldown_cache_key(model_id: str) -> str: return "deployment:" + model_id + ":cooldown" + def _corrected_active_cooldown( + self, + key: str, + result: Mapping[str, Any], + current_time: float, + ) -> CooldownCacheValue | None: + """ + Return a CooldownCacheValue if the cooldown is still active, or None if it has expired. + + Also corrects the in-memory TTL when DualCache promotes a Redis entry using the + default 600s TTL instead of the true remaining cooldown time. + """ + 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) + return None + current_expiry: Final = self.cache.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) + return cooldown_cache_value + async def async_get_active_cooldowns( self, model_ids: list[str], parent_otel_span: Span | None ) -> list[tuple[str, CooldownCacheValue]]: @@ -117,11 +148,13 @@ class CooldownCache: if results is None or all(v is None for v in results): return active_cooldowns - # Process the results + current_time: Final = time.time() for model_id, result in zip(model_ids, results): if result and isinstance(result, dict): - cooldown_cache_value = CooldownCacheValue(**result) - active_cooldowns.append((model_id, cooldown_cache_value)) + key = CooldownCache.get_cooldown_cache_key(model_id) + cooldown_cache_value = self._corrected_active_cooldown(key, result, current_time) + if cooldown_cache_value is not None: + active_cooldowns.append((model_id, cooldown_cache_value)) return active_cooldowns @@ -134,11 +167,13 @@ class CooldownCache: results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or [] active_cooldowns: Final = [] - # Process the results + current_time: Final = time.time() for model_id, result in zip(model_ids, results): if result and isinstance(result, dict): - cooldown_cache_value = CooldownCacheValue(**result) - active_cooldowns.append((model_id, cooldown_cache_value)) + key = CooldownCache.get_cooldown_cache_key(model_id) + cooldown_cache_value = self._corrected_active_cooldown(key, result, current_time) + if cooldown_cache_value is not None: + active_cooldowns.append((model_id, cooldown_cache_value)) return active_cooldowns diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 2b26928a21c..39618a6f182 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -8,6 +8,8 @@ Router cooldown handlers import asyncio import math +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import litellm @@ -58,6 +60,148 @@ def is_advisor_orchestration_failure(exception: BaseException | None) -> bool: return bool(getattr(exception, _ADVISOR_ORCHESTRATION_FAILURE_ATTR, False)) +_EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = ( + # ContentPolicyViolationError subclasses BadRequestError, so it must be checked first. + (litellm.ContentPolicyViolationError, "ContentPolicyViolationErrorAllowedFails"), + (litellm.BadRequestError, "BadRequestErrorAllowedFails"), + (litellm.AuthenticationError, "AuthenticationErrorAllowedFails"), + (litellm.Timeout, "TimeoutErrorAllowedFails"), + (litellm.RateLimitError, "RateLimitErrorAllowedFails"), + (litellm.InternalServerError, "InternalServerErrorAllowedFails"), + (litellm.ServiceUnavailableError, "ServiceUnavailableErrorAllowedFails"), + (litellm.BadGatewayError, "BadGatewayErrorAllowedFails"), + (litellm.NotFoundError, "NotFoundErrorAllowedFails"), +) + + +def _first_present(*sources: Mapping[str, Any] | None, key: str) -> int | float | None: + """Return *key* from the first source mapping where it's set, so callers can + support a setting living in more than one deployment config location. Sources + are checked in order from most to least specific to that setting.""" + for source in sources: + if source is None: + continue + value = source.get(key) + if value is not None: + return value + return None + + +def _get_deployment_cooldown_policy( + litellm_router_instance: LitellmRouter, + deployment: str, +) -> tuple[Mapping[str, int] | None, int | None]: + """Return (allowed_fails_policy, allowed_fails) from deployment model_info, or (None, None). + + `model_info` is the only supported location for these two fields (unlike + `cooldown_time`, they have no pre-existing `litellm_params` precedent): `litellm_params` + gets copied wholesale into the actual provider call kwargs (see e.g. + Router._image_generation's `data = deployment["litellm_params"].copy()`), so a new + field placed there would leak into the outgoing LLM request instead of staying + router-internal. + """ + dep: Final = litellm_router_instance.get_model_info(id=deployment) + if dep is None: + return None, None + mi: Final[Mapping[str, Any]] = dep.get("model_info") or MappingProxyType({}) + raw: Final = mi.get("allowed_fails_policy") + policy: Final[Mapping[str, int] | None] = raw if isinstance(raw, dict) else None + allowed: Final[int | None] = mi.get("allowed_fails") + return policy, allowed + + +def _resolve_allowed_fails_from_policy( + policy: Mapping[str, int] | None, + exception: Exception, +) -> int | None: + """Match *exception* against *policy* and return the configured allowed-fail count, or None.""" + if policy is None: + return None + for exc_type, field in _EXCEPTION_POLICY_FIELDS: + if isinstance(exception, exc_type): + value = policy.get(field) + if value is not None: + return value + return None + + +def _should_cooldown_based_on_deployment_policy( + litellm_router_instance: LitellmRouter, + deployment: str, + original_exception: Exception, + dep_policy: Mapping[str, int] | None, + dep_allowed_fails: int | None, + is_single_deployment_model_group: bool, +) -> bool: + """Resolve deployment-level allowed-fails and delegate to the shared counting logic. + + When the deployment's policy doesn't cover *original_exception*'s type and no + deployment-wide `allowed_fails` is set either, defer to router-level behavior + instead of forcing an immediate cooldown. + + A generic, deployment-wide `allowed_fails` predates this feature's per-exception-type + policy and is a much less deliberate opt-in, so on a single-deployment model group it + still defers to the "avoid cooldowns on single deployment model groups" safety net + (see `_should_cooldown_deployment`'s BASE CASE) rather than silently disabling it. An + explicit, named-exception-type `allowed_fails_policy` entry is unambiguous enough to + override that safety net, matching `_has_explicit_allowed_fails_policy_for_exception`. + """ + allowed_fails_from_policy: Final = _resolve_allowed_fails_from_policy(dep_policy, original_exception) + if allowed_fails_from_policy is None and dep_allowed_fails is not None and is_single_deployment_model_group: + return False + + allowed_fails_override: Final[int | None] = ( + allowed_fails_from_policy if allowed_fails_from_policy is not None else dep_allowed_fails + ) + cache_key_suffix: Final[str | None] = ( + type(original_exception).__name__ + if allowed_fails_from_policy is not None + else ("generic" if dep_allowed_fails is not None else None) + ) + + dep: Final = litellm_router_instance.get_model_info(id=deployment) + cooldown_time_override: Final = ( + _first_present(dep.get("model_info"), dep.get("litellm_params"), key="cooldown_time") + if dep is not None + else None + ) + + return should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=litellm_router_instance, + deployment=deployment, + original_exception=original_exception, + allowed_fails_override=allowed_fails_override, + cooldown_time_override=cooldown_time_override, + cache_key_suffix=cache_key_suffix, + ) + + +def _has_explicit_allowed_fails_policy_for_exception( + litellm_router_instance: LitellmRouter, + deployment: str | None, + original_exception: Exception, +) -> bool: + """True if this deployment has an explicit, deployment-level allowed_fails_policy + entry matching *original_exception*'s type. + + `_is_cooldown_required` skips cooldown evaluation for most 4XX errors (BadRequestError, + ContentPolicyViolationError) by default, since a generic client error is usually not the + deployment's fault. A deployment-level allowed_fails_policy entry naming that exact + exception type is this PR's own per-deployment opt-in, so it overrides that default. + + Deliberately scoped to the deployment level only, and to the named-exception-type + policy dict rather than a plain `allowed_fails` integer: a pre-existing router-wide + `allowed_fails_policy` (or a deployment's generic `allowed_fails`) predates this + feature and must keep its existing behavior for 4XX types `_is_cooldown_required` + already excludes, rather than silently start cooling down deployments whose configs + never opted into this specific override. + """ + if deployment is None: + return False + dep_policy, _ = _get_deployment_cooldown_policy(litellm_router_instance, deployment) + return _resolve_allowed_fails_from_policy(dep_policy, original_exception) is not None + + def _is_cooldown_required( litellm_router_instance: LitellmRouter, model_id: str, @@ -155,6 +299,10 @@ def _should_run_cooldown_logic( model_id=deployment, exception_status=exception_status, exception_str=str(original_exception), + ) and not _has_explicit_allowed_fails_policy_for_exception( + litellm_router_instance=litellm_router_instance, + deployment=deployment, + original_exception=original_exception, ): verbose_router_logger.debug("Should Not Run Cooldown Logic: _is_cooldown_required returned False") return False @@ -190,11 +338,24 @@ def _should_cooldown_deployment( - v1 logic (Legacy): if allowed fails or allowed fail policy set, coolsdown if num fails in this minute > allowed fails """ - ## BASE CASE - single deployment model_group: Final = litellm_router_instance.get_model_group(id=deployment) is_single_deployment_model_group = False if model_group is not None and len(model_group) == 1: is_single_deployment_model_group = True + + ## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level) + dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(litellm_router_instance, deployment) + if dep_policy is not None or dep_allowed_fails is not None: + return _should_cooldown_based_on_deployment_policy( + litellm_router_instance, + deployment, + original_exception, + dep_policy, + dep_allowed_fails, + is_single_deployment_model_group, + ) + + ## BASE CASE - single deployment if ( litellm_router_instance.allowed_fails_policy is None and _is_allowed_fails_set_on_router(litellm_router_instance=litellm_router_instance) is False @@ -382,29 +543,50 @@ def should_cooldown_based_on_allowed_fails_policy( litellm_router_instance: LitellmRouter, deployment: str, original_exception: Any, + allowed_fails_override: int | None = None, + cooldown_time_override: float | None = None, + cache_key_suffix: str | None = None, ) -> bool: """ Check if fails are within the allowed limit and update the number of fails. + When *allowed_fails_override* / *cooldown_time_override* are supplied they + take precedence over the router-level values (used by deployment-level overrides). + + 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. + Returns: - True if fails exceed the allowed limit (should cooldown) - False if fails are within the allowed limit (should not cooldown) """ - allowed_fails: Final = ( - litellm_router_instance.get_allowed_fails_from_policy( - exception=original_exception, - ) - or litellm_router_instance.allowed_fails + allowed_fails_from_policy: Final = litellm_router_instance.get_allowed_fails_from_policy( + exception=original_exception + ) + allowed_fails: Final = ( + allowed_fails_override + if allowed_fails_override is not None + else ( + allowed_fails_from_policy + if allowed_fails_from_policy is not None + else litellm_router_instance.allowed_fails + ) + ) + cooldown_time: Final = ( + cooldown_time_override + if cooldown_time_override is not None + else (litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS) ) - cooldown_time: Final = litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS - current_fails: Final = litellm_router_instance.failed_calls.get_cache(key=deployment) or 0 + 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 if updated_fails > allowed_fails: return True else: - litellm_router_instance.failed_calls.set_cache(key=deployment, value=updated_fails, ttl=cooldown_time) + litellm_router_instance.failed_calls.set_cache(key=cache_key, value=updated_fails, ttl=cooldown_time) return False diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index ef48ccc821e..63bc5203417 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -1,5 +1,6 @@ import hashlib import json +from collections.abc import Mapping from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Any, Final @@ -12,6 +13,16 @@ from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, get_fallback_error_info, ) +from litellm.router_utils.batch_utils import _get_router_metadata_variable_name +from litellm.router_utils.cooldown_handlers import ( + _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils + _set_cooldown_deployments, # pyright: ignore[reportPrivateUsage] - shared helper, used across router_utils + cast_exception_status_to_int, + is_advisor_orchestration_failure, +) +from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + increment_deployment_failures_for_current_minute, +) from litellm.types.router import LiteLLMParamsTypedDict if TYPE_CHECKING: @@ -21,6 +32,116 @@ if TYPE_CHECKING: else: LitellmRouter = Any +# Status codes a generic API call's caller-supplied resource id can trigger on its own +# (e.g. a nonexistent file/batch/thread id), independent of the selected deployment's health. +_REQUEST_SCOPED_STATUS_CODES: Final = frozenset((404,)) + + +def _trigger_cooldown_for_failed_deployment( + litellm_router: LitellmRouter, + kwargs: Mapping[str, Any], + exception: Exception, +) -> None: + """ + Trigger cooldown for a failed fallback deployment. + + In the fallback path the normal failure-callback cooldown is skipped because the + Logging object sets has_logged_async_failure=True after the first failure and + blocks all subsequent failure callbacks. This helper ensures every failed + fallback deployment is evaluated for cooldown regardless. + """ + try: + if is_advisor_orchestration_failure(exception): + verbose_router_logger.debug( + "Not triggering cooldown for fallback deployment: failure originated " + "from advisor orchestration, not the selected deployment." + ) + return + + exception_status: Final[str | int] = getattr(exception, "status_code", "") + + # Generic API calls (files, batches, threads, rerank, ...) take a caller-supplied + # resource id, so a 404 there usually means "that id doesn't exist" rather than + # "this deployment is unhealthy". Left unguarded, one bad id would 404 every + # deployment in the fallback chain and cool all of them down from a single request. + if ( + kwargs.get("original_generic_function") is not None + and cast_exception_status_to_int(exception_status) in _REQUEST_SCOPED_STATUS_CODES + ): + verbose_router_logger.debug( + "Not triggering cooldown for fallback deployment: status %s on a generic API " + "call is caller-attributable, not a deployment health signal.", + exception_status, + ) + return + + # The proxy's `x-litellm-timeout` header lets a caller set an arbitrarily short + # timeout, which litellm.Timeout reports as status 408 regardless of the deployment's + # actual health. Left unguarded, a caller could force a 408 on every deployment in + # the fallback chain from a single request with a near-zero timeout. + if kwargs.get("client_side_timeout") and cast_exception_status_to_int(exception_status) == 408: + verbose_router_logger.debug( + "Not triggering cooldown for fallback deployment: a caller-supplied " + "x-litellm-timeout caused this 408, not deployment health." + ) + return + + # Only Router._set_failed_deployment_id_on_exception()'s server-stamped id is + # trusted here: a metadata-bucket lookup (e.g. "metadata"/"litellm_metadata") + # can't reliably tell a caller-supplied bucket from a router-authored one + # without knowing this call's function_name, so a client with permission to + # set metadata could otherwise get an arbitrary deployment cooled down. + deployment_id: Final[str | None] = getattr(exception, "failed_deployment_id", None) + + if deployment_id is None: + verbose_router_logger.debug("Cannot trigger cooldown for fallback: no failed_deployment_id on exception") + return + + # Priority: deployment config > response header > router default, matching + # Router.deployment_callback_on_failure's precedence for the primary path. + deployment_dict: Final = litellm_router.get_model_info(id=deployment_id) + deployment_cooldown: Final = ( + _first_present( + deployment_dict.get("model_info"), deployment_dict.get("litellm_params"), key="cooldown_time" + ) + if deployment_dict is not None + else None + ) + exception_headers: Final = litellm.litellm_core_utils.exception_mapping_utils._get_response_headers( + original_exception=exception + ) + _get_retry_after: Final = ( + litellm.utils._get_retry_after_from_exception_header # pyright: ignore[reportPrivateUsage] - as router.py + ) + header_cooldown: Final = ( + _get_retry_after(response_headers=exception_headers) if exception_headers is not None else None + ) + time_to_cooldown: Final = ( + deployment_cooldown + if deployment_cooldown is not None and deployment_cooldown >= 0 + else ( + header_cooldown + if header_cooldown is not None and header_cooldown >= 0 + else litellm_router.cooldown_time + ) + ) + + increment_deployment_failures_for_current_minute( + litellm_router_instance=litellm_router, + deployment_id=deployment_id, + ) + _set_cooldown_deployments( + litellm_router_instance=litellm_router, + exception_status=exception_status, + original_exception=exception, + deployment=deployment_id, + time_to_cooldown=time_to_cooldown, + ) + + verbose_router_logger.debug("Triggered cooldown for fallback deployment %s", deployment_id) + except Exception as e: # noqa: BLE001 - best-effort cooldown trigger must never break the fallback response itself + verbose_router_logger.debug("Error triggering cooldown for fallback deployment: %s", e) + def fallback_attempt_key(fallback_target: object) -> str | None: """ @@ -131,6 +252,28 @@ def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[li return fallback_model_group, generic_fallback_idx +PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file") + + +def _get_fallback_target_model_group(fallback_entry: str | Mapping[str, object]) -> str | None: + if isinstance(fallback_entry, str): + return fallback_entry + target: Final = fallback_entry.get("model") + return target if isinstance(target, str) else None + + +def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: + """ + True when the request names a file that only exists under one provider's credentials. + + Batch and fine-tuning jobs are created from a file the caller already uploaded, and + that file lives in the account of the deployment that stored it. Handing the id to a + different model group can only fail, and the second provider's error replaces the + error the caller actually needs to see. + """ + return any(kwargs.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS) + + async def run_async_fallback( *args: tuple[Any], litellm_router: LitellmRouter, @@ -176,6 +319,10 @@ async def run_async_fallback( error_from_fallbacks = original_exception fallback_errors = (get_fallback_error_info(original_exception),) + metadata_variable_name: Final = _get_router_metadata_variable_name( + function_name=getattr(kwargs.get("original_function"), "__name__", None) + ) + same_model_group_only: Final = references_provider_scoped_resource(kwargs) # Read out of kwargs and narrowed here rather than declared as a parameter: every caller # reaches this function by spreading a loosely-typed kwargs dict, so a declared parameter # would carry an annotation that no call site can actually be checked against. @@ -188,6 +335,13 @@ async def run_async_fallback( for mg in fallback_model_group: if mg == original_model_group: continue + if same_model_group_only and _get_fallback_target_model_group(mg) != original_model_group: + verbose_router_logger.info( + "Skipping fallback to model_group = %s: request is pinned to model_group = %s by its uploaded file", + mask_sensitive_structure(mg), + original_model_group, + ) + continue attempt_key = fallback_attempt_key(mg) if attempt_key is not None: if attempt_key in attempted: @@ -205,9 +359,10 @@ async def run_async_fallback( kwargs["model"] = mg elif isinstance(mg, dict): kwargs.update(mg) - kwargs.setdefault("metadata", {}).update( - {"model_group": kwargs.get("model", None)} - ) # update model_group used, if fallbacks are done + kwargs[metadata_variable_name] = { + **(kwargs.get(metadata_variable_name) or {}), + "model_group": kwargs.get("model", None), + } fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks @@ -236,6 +391,13 @@ async def run_async_fallback( kwargs=kwargs, original_exception=original_exception, ) + logging_obj = kwargs.get("litellm_logging_obj") + if logging_obj is not None and logging_obj.model_call_details.get("has_logged_async_failure", False): + _trigger_cooldown_for_failed_deployment( + litellm_router=litellm_router, + kwargs=kwargs, + exception=e, + ) raise error_from_fallbacks diff --git a/litellm/types/memory_management.py b/litellm/types/memory_management.py index 04a2a0c1905..153de0c6cb9 100644 --- a/litellm/types/memory_management.py +++ b/litellm/types/memory_management.py @@ -3,7 +3,6 @@ Pydantic models for Memory management endpoints. """ from datetime import datetime -from typing import Any from pydantic import BaseModel, Field @@ -12,7 +11,7 @@ class LiteLLM_MemoryRow(BaseModel): memory_id: str key: str value: str - metadata: Any | None = None + metadata: object | None = None user_id: str | None = None team_id: str | None = None created_at: datetime | None = None @@ -24,7 +23,7 @@ class LiteLLM_MemoryRow(BaseModel): class MemoryCreateRequest(BaseModel): key: str = Field(..., description="Memory key (acts as the namespace in the URL).") value: str = Field(..., description="Memory content. Typically markdown/text for LLM context.") - metadata: Any | None = Field( + metadata: object | None = Field( default=None, description="Optional JSON metadata (tags, structured fields).", ) @@ -40,7 +39,7 @@ class MemoryCreateRequest(BaseModel): class MemoryUpdateRequest(BaseModel): value: str | None = None - metadata: Any | None = None + metadata: object | None = None # Only honored on create (when the row doesn't yet exist) and only for # PROXY_ADMIN callers — mirrors MemoryCreateRequest so admins can bootstrap # rows scoped to another user/team via PUT, not just POST. diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index fc1e6d15fd4..16d08b33150 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -18,6 +18,7 @@ class GroupByDimension(str, Enum): class SpendMetrics(BaseModel): spend: float = Field(default=0.0) + flat_cost: float = Field(default=0.0) prompt_tokens: int = Field(default=0) completion_tokens: int = Field(default=0) cache_read_input_tokens: int = Field(default=0) @@ -75,6 +76,7 @@ class DailySpendData(BaseModel): class DailySpendMetadata(BaseModel): total_spend: float = Field(default=0.0) + total_flat_cost: float = Field(default=0.0) total_prompt_tokens: int = Field(default=0) total_completion_tokens: int = Field(default=0) total_tokens: int = Field(default=0) diff --git a/litellm/types/router.py b/litellm/types/router.py index e166d844735..4f8c133c20b 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -5,7 +5,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum from dataclasses import dataclass -from typing import Any, Final, Generic, Literal, TypeVar, get_type_hints +from typing import Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator @@ -127,6 +127,14 @@ class UpdateRouterConfig(BaseModel): model_config = ConfigDict(protected_namespaces=()) +def _as_utc(value: datetime.datetime | None) -> datetime.datetime | None: + if value is None: + return None + if value.tzinfo is None: + return value.replace(tzinfo=datetime.timezone.utc) + return value.astimezone(datetime.timezone.utc) + + class ModelInfo(MirroredPricingParams): id: str | None # Allow id to be optional on input, but it will always be present as a str in the model instance db_model: bool = False # used for proxy - to separate models which are stored in the db vs. config. @@ -151,6 +159,17 @@ class ModelInfo(MirroredPricingParams): # admin-toggled pause flag; mirrors LiteLLM_ProxyModelTable.blocked blocked: bool | None = None + # Bounds live on the model rather than litellm.constants: names there reach + # litellm/__init__ through several modules' star re-exports, and a Final rebound that + # way trips the basedpyright gate. + MAX_PTU_COUNT: ClassVar[int] = 1_000_000 + MAX_COST_PER_PTU_PER_HOUR: ClassVar[float] = 1_000_000.0 + + ptu_count: int | None = None + cost_per_ptu_per_hour: float | None = None + ptu_effective_from: datetime.datetime | None = None + ptu_effective_to: datetime.datetime | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided @@ -158,6 +177,23 @@ class ModelInfo(MirroredPricingParams): id = str(id) super().__init__(id=id, **params) + @model_validator(mode="after") + def _validate_ptu_bounds(self) -> "ModelInfo": + if self.ptu_count is not None and not 0 < self.ptu_count <= self.MAX_PTU_COUNT: + raise ValueError(f"ptu_count must be a positive integer no greater than {self.MAX_PTU_COUNT}") + if ( + self.cost_per_ptu_per_hour is not None + and not 0 <= self.cost_per_ptu_per_hour <= self.MAX_COST_PER_PTU_PER_HOUR + ): + raise ValueError( + f"cost_per_ptu_per_hour must be a finite number between 0 and {self.MAX_COST_PER_PTU_PER_HOUR}" + ) + start: Final = _as_utc(self.ptu_effective_from) + end: Final = _as_utc(self.ptu_effective_to) + if start is not None and end is not None and end <= start: + raise ValueError("ptu_effective_to must be after ptu_effective_from") + return self + model_config = ConfigDict(extra="allow") def __contains__(self, key) -> bool: @@ -245,6 +281,12 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): # Deployment budgets max_budget: float | None = None budget_duration: str | None = None + keepalive_seconds: float | None = None + # keepalive_seconds is operator-only by default: a client's request-level + # value is ignored unless the deployment opts in here. Prevents a client + # from unilaterally enabling heartbeats (and the LB-idle-timeout evasion + # that comes with them) for a deployment that never configured them. + allow_client_keepalive_override: bool | None = False use_in_pass_through: bool | None = False use_litellm_proxy: bool | None = False use_chat_completions_api: bool | None = None @@ -421,6 +463,11 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): # deployment budgets max_budget: float | None budget_duration: str | None + keepalive_seconds: float | None + allow_client_keepalive_override: bool | None + + # per-deployment cooldown override + cooldown_time: float | None class DeploymentTypedDict(TypedDict, total=False): @@ -513,6 +560,9 @@ class AllowedFailsPolicy(BaseModel): RateLimitErrorAllowedFails: int | None = None ContentPolicyViolationErrorAllowedFails: int | None = None InternalServerErrorAllowedFails: int | None = None + ServiceUnavailableErrorAllowedFails: int | None = None + BadGatewayErrorAllowedFails: int | None = None + NotFoundErrorAllowedFails: int | None = None class AlertingConfig(BaseModel): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8f34e23a5ae..c9bb7bb1dc3 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3129,6 +3129,7 @@ class StandardAuditLogPayload(TypedDict): class StandardLoggingPayload(TypedDict): id: str trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries) + session_id: str # End-user/conversation session id (litellm_session_id), independent of trace_id litellm_call_id: str | None # UUID returned in x-litellm-call-id response header call_type: str stream: bool | None @@ -3384,6 +3385,8 @@ agentic_loop_internal_litellm_params: Final = [ "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_session_scoped", "_code_interpreter_interception_converted_stream", + "_websearch_interception_emit_native_blocks", + "_websearch_interception_converted_stream", ] # Proxy-owned callback credentials, stamped from admin-configured team/key callback @@ -3398,6 +3401,8 @@ all_litellm_params = ( + [ "metadata", "litellm_metadata", + "keepalive_seconds", + "allow_client_keepalive_override", "litellm_trace_id", "litellm_request_debug", "guardrails", diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index d1d4a39da1e..474c652ff3a 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -277,6 +277,11 @@ class LiteLLM_ManagedVectorStoreIndex(BaseModel): updated_by: str | None = None +class IndexListResponse(BaseModel): + object: Literal["list"] = "list" + data: tuple[LiteLLM_ManagedVectorStoreIndex, ...] + + class VectorStoreIndexType(str, Enum): """Type of vector store index""" diff --git a/litellm/utils.py b/litellm/utils.py index 911de83b785..87937c99a0c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -711,14 +711,71 @@ def _remove_thought_signatures_from_messages(messages: list, thought_signature_s return processed_messages +def _restore_correlation_context_if_supported(logging_obj: object) -> None: + """Call logging_obj._restore_correlation_context() if it's actually there. + + Some call sites (tests, narrow unit paths) inject a minimal stand-in + object as litellm_logging_obj instead of a real Logging instance - this + method is new plumbing specific to request_correlation_in_logs, not part + of any pre-existing stand-in's expected interface. `object` (not `Any`) + is deliberate: the getattr() below is exactly how this stays type-safe + while still tolerating a stand-in that lacks the method. + """ + restore: Final = getattr(logging_obj, "_restore_correlation_context", None) + if restore is not None: + restore() + + +def _is_streaming_response_for_correlation(result: object) -> bool: + """True if `result` is a lazy stream wrapper rather than an already-complete response. + + Only wrapper_async() consults this - it must NOT restore the originating + Task's trace_id/session_id as soon as a streaming call returns this: the + caller is about to iterate it over however many subsequent lines of their + own code, and those log lines should still show this call's ids, not the + pre-call ones. This is safe specifically because each async call already + runs in its own asyncio Task with its own copy of the contextvars, so + leaving it "open" can only affect that one Task, never a different, + unrelated future request - Tasks, unlike a thread pool's worker threads, + are never recycled across requests. The corresponding terminal handler + (async_success_handler, dispatched once the full stream is actually + assembled) is what restores it once streaming genuinely finishes. + + wrapper() (the sync path) does NOT consult this at all: sync calls pass + supports_correlation_logging=False into function_setup()/Logging(), so + they never stamp trace_id/session_id in the first place - a plain OS + thread has no per-call isolation the way an asyncio Task does, and a + thread pool's worker threads *are* recycled across unrelated requests, so + stamping ids there without a safe restore mechanism could permanently + misattribute a later, unrelated request's logs. Full sync support is + deferred to a follow-up PR with its own restore mechanism; see + Logging.__init__'s supports_correlation_logging parameter. + + Genuinely circular otherwise: utils.py -> streaming_handler.py -> + redact_messages.py -> llms/vertex_ai/common_utils.py -> utils.py, which + needs names (supports_response_schema, etc.) this module hasn't finished + defining yet at that point in its own top-to-bottom execution. + """ + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + return isinstance(result, CustomStreamWrapper) + + +# Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc. def function_setup( - original_function: str, rules_obj, start_time, *args, **kwargs -): # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc. + original_function: str, + rules_obj: Rules, + start_time: datetime.datetime, + *args: Any, # positional passthrough to the wrapped LLM call (ANN401 ignored, see ruff-strict.toml) + is_async_call: bool = True, + **kwargs: Any, # kwargs-ok: forwarded to Logging()/callbacks, varies per call_type +) -> tuple[LiteLLMLoggingObject, dict[str, Any]]: ### NOTICES ### if litellm.set_verbose is True: verbose_logger.warning( "`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs." ) + logging_obj: LiteLLMLoggingObject | None = None # rebind-ok: set to the real object further down on success try: global callback_list, add_breadcrumb, user_logger_fn, Logging @@ -1001,7 +1058,8 @@ def function_setup( ): stream = True get_litellm_logging_class: Final = getattr(sys.modules[__name__], "get_litellm_logging_class") - logging_obj: Final = get_litellm_logging_class()( # Victim for object pool + # Victim for object pool + logging_obj = get_litellm_logging_class()( # rebind-ok: 2nd assignment to logging_obj (see initial None above) model=model, messages=messages, stream=stream, @@ -1016,6 +1074,7 @@ def function_setup( dynamic_async_failure_callbacks=dynamic_async_failure_callbacks, kwargs=kwargs, applied_guardrails=applied_guardrails, + supports_correlation_logging=is_async_call, ) ## check if metadata is passed in @@ -1040,6 +1099,15 @@ def function_setup( ) return logging_obj, kwargs except Exception as e: + # If Logging() was constructed above before this failed, its __init__ already + # mutated trace_id_var/session_id_var - restore them *before* logging the + # exception below, since we're about to raise without ever returning + # logging_obj to the caller's wrapper()/wrapper_async() (which would + # otherwise be the one doing this restore). Restoring first means this + # diagnostic log line itself doesn't get stamped with a call's ids when + # that call never actually produced a usable logging object. + if logging_obj is not None: + _restore_correlation_context_if_supported(logging_obj) verbose_logger.exception("litellm.utils.py::function_setup() - [Non-Blocking] Error in function_setup") raise e @@ -1296,7 +1364,9 @@ def client(original_function): try: if logging_obj is None: - logging_obj, kwargs = function_setup(original_function.__name__, rules_obj, start_time, *args, **kwargs) + logging_obj, kwargs = function_setup( + original_function.__name__, rules_obj, start_time, *args, is_async_call=False, **kwargs + ) # Type assertion: logging_obj is guaranteed to be non-None after function_setup assert logging_obj is not None, "logging_obj should not be None after function_setup" @@ -1807,9 +1877,11 @@ def client(original_function): kwargs["retry_strategy"] = "exponential_backoff_retry" elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" - return await litellm.acompletion_with_retries(*args, **kwargs) + result = await litellm.acompletion_with_retries(*args, **kwargs) except Exception: pass + else: + return result elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict @@ -1820,7 +1892,8 @@ def client(original_function): args[0] = context_window_fallback_dict[model] else: kwargs["model"] = context_window_fallback_dict[model] - return await original_function(*args, **kwargs) + result = await original_function(*args, **kwargs) + return result elif call_type == CallTypes.aresponses.value: _is_litellm_router_call = "model_group" in ( kwargs.get("metadata") or {} @@ -1837,9 +1910,11 @@ def client(original_function): kwargs["retry_strategy"] = "exponential_backoff_retry" elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" - return await litellm.aresponses_with_retries(*args, **kwargs) + result = await litellm.aresponses_with_retries(*args, **kwargs) except Exception: pass + else: + return result deployment_num_retries: Final = kwargs.get("num_retries") if deployment_num_retries is not None: @@ -1849,6 +1924,21 @@ def client(original_function): setattr(e, "timeout", timeout) raise e + finally: + # Restore trace_id/session_id contextvars to their pre-call value once + # this call (in this asyncio Task) is fully done - see + # request_correlation_in_logs. Unlike wrapper()'s sync path, it's safe to + # skip restoring when returning a stream: each async call already runs in + # its own Task with its own copy of the contextvars (asyncio.create_task + # copies context at creation), so leaving this Task's own view "open" + # while the caller iterates the stream can only affect that one Task - + # never a different, unrelated future request, since Tasks (unlike a + # thread pool's worker threads) are never recycled across requests. The + # corresponding terminal handler (async_success_handler) restores it once + # streaming genuinely finishes; aclose()/__del__ cover early termination. + if not _is_streaming_response_for_correlation(result): + _restore_correlation_context_if_supported(logging_obj) + get_coroutine_checker: Final = getattr(sys.modules[__name__], "get_coroutine_checker") is_coroutine: Final = get_coroutine_checker().is_async_callable(original_function) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8982b4f2565..951c114b0a9 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -2808,6 +2808,7 @@ }, "azure/codex-mini": { "cache_read_input_token_cost": 3.75e-07, + "deprecation_date": "2026-11-15", "input_cost_per_token": 1.5e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -3627,7 +3628,7 @@ "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, "azure/eu/gpt-4o-2024-08-06": { - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.375e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -3644,7 +3645,7 @@ "supports_vision": true }, "azure/eu/gpt-4o-2024-11-20": { - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "cache_creation_input_token_cost": 1.38e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -3661,6 +3662,7 @@ }, "azure/eu/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 8.3e-08, + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -3742,6 +3744,7 @@ }, "azure/eu/gpt-5-2025-08-07": { "cache_read_input_token_cost": 1.375e-07, + "deprecation_date": "2027-02-09", "input_cost_per_token": 1.375e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -3774,6 +3777,7 @@ }, "azure/eu/gpt-5-mini-2025-08-07": { "cache_read_input_token_cost": 2.75e-08, + "deprecation_date": "2027-02-09", "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -3840,6 +3844,7 @@ }, "azure/eu/gpt-5.1-chat": { "cache_read_input_token_cost": 1.4e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -3934,6 +3939,7 @@ }, "azure/eu/gpt-5-nano-2025-08-07": { "cache_read_input_token_cost": 5.5e-09, + "deprecation_date": "2027-02-09", "input_cost_per_token": 5.5e-08, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -3966,6 +3972,7 @@ }, "azure/eu/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -4011,6 +4018,7 @@ }, "azure/eu/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, + "deprecation_date": "2026-10-01", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -4027,7 +4035,7 @@ }, "azure/global-standard/gpt-4o-2024-08-06": { "cache_read_input_token_cost": 1.25e-06, - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4044,7 +4052,7 @@ }, "azure/global-standard/gpt-4o-2024-11-20": { "cache_read_input_token_cost": 1.25e-06, - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4073,7 +4081,7 @@ "supports_vision": true }, "azure/global/gpt-4o-2024-08-06": { - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4090,7 +4098,7 @@ "supports_vision": true }, "azure/global/gpt-4o-2024-11-20": { - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4142,6 +4150,7 @@ }, "azure/global/gpt-5.1-chat": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4476,7 +4485,7 @@ "supports_web_search": false }, "azure/gpt-4.1-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -4543,7 +4552,7 @@ "supports_web_search": false }, "azure/gpt-4.1-mini-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 4e-07, "input_cost_per_token_batches": 2e-07, @@ -4609,7 +4618,7 @@ "supports_vision": true }, "azure/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, @@ -4677,6 +4686,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-05-13": { + "deprecation_date": "2026-10-01", "input_cost_per_token": 5e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4691,7 +4701,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-08-06": { - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4708,7 +4718,7 @@ "supports_vision": true }, "azure/gpt-4o-2024-11-20": { - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.25e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -4725,6 +4735,7 @@ "supports_vision": true }, "azure/gpt-audio-2025-08-28": { + "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4756,6 +4767,7 @@ "supports_vision": false }, "azure/gpt-audio-1.5-2026-02-23": { + "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -4787,6 +4799,7 @@ "supports_vision": false }, "azure/gpt-audio-mini-2025-10-06": { + "deprecation_date": "2027-04-06", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "azure", @@ -4866,6 +4879,7 @@ }, "azure/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -4933,6 +4947,7 @@ "azure/gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_token_cost": 4e-06, + "deprecation_date": "2027-03-02", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, @@ -4965,6 +4980,7 @@ "azure/gpt-realtime-1.5-2026-02-23": { "cache_creation_input_audio_token_cost": 4e-06, "cache_read_input_token_cost": 4e-06, + "deprecation_date": "2027-08-24", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, @@ -5102,6 +5118,7 @@ "supports_tool_choice": true }, "azure/gpt-4o-transcribe": { + "deprecation_date": "2026-10-15", "input_cost_per_audio_token": 2.5e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -5114,6 +5131,7 @@ ] }, "azure/gpt-4o-transcribe-diarize": { + "deprecation_date": "2027-04-15", "input_cost_per_audio_token": 2.5e-06, "input_cost_per_token": 2.5e-06, "litellm_provider": "azure", @@ -5145,6 +5163,7 @@ "azure/gpt-5.1-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2027-05-15", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "azure", @@ -5182,6 +5201,7 @@ "azure/gpt-5.1-chat-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "azure", @@ -5218,6 +5238,7 @@ "azure/gpt-5.1-codex-2025-11-13": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2027-05-15", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "azure", @@ -5251,6 +5272,7 @@ "azure/gpt-5.1-codex-mini-2025-11-13": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2027-05-15", "input_cost_per_token": 2.5e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "azure", @@ -5315,6 +5337,7 @@ }, "azure/gpt-5-2025-08-07": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2027-02-09", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5347,6 +5370,7 @@ }, "azure/gpt-5-chat": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -5380,6 +5404,7 @@ }, "azure/gpt-5-chat-latest": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -5412,6 +5437,7 @@ }, "azure/gpt-5-codex": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2027-03-17", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5474,6 +5500,7 @@ }, "azure/gpt-5-mini-2025-08-07": { "cache_read_input_token_cost": 2.5e-08, + "deprecation_date": "2027-02-09", "input_cost_per_token": 2.5e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5538,6 +5565,7 @@ }, "azure/gpt-5-nano-2025-08-07": { "cache_read_input_token_cost": 5e-09, + "deprecation_date": "2027-02-09", "input_cost_per_token": 5e-08, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5569,6 +5597,7 @@ "supports_vision": true }, "azure/gpt-5-pro": { + "deprecation_date": "2027-04-07", "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5633,6 +5662,7 @@ }, "azure/gpt-5.1-chat": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -5697,6 +5727,7 @@ }, "azure/gpt-5.1-codex-max": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2027-05-18", "input_cost_per_token": 1.25e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5791,6 +5822,7 @@ "azure/gpt-5.2-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2027-06-08", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "azure", @@ -5827,6 +5859,7 @@ "azure/gpt-5.2-chat": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "azure", @@ -5861,6 +5894,7 @@ "azure/gpt-5.2-chat-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-05-13", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "azure", @@ -5894,6 +5928,7 @@ }, "azure/gpt-5.2-codex": { "cache_read_input_token_cost": 1.75e-07, + "deprecation_date": "2027-07-13", "input_cost_per_token": 1.75e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -5925,6 +5960,7 @@ "azure/gpt-5.3-chat": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "azure", @@ -5958,6 +5994,7 @@ }, "azure/gpt-5.3-codex": { "cache_read_input_token_cost": 1.75e-07, + "deprecation_date": "2027-08-24", "input_cost_per_token": 1.75e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -6164,6 +6201,7 @@ "cache_read_input_token_cost_above_272k_tokens": 5e-07, "cache_read_input_token_cost_priority": 5e-07, "cache_read_input_token_cost_above_272k_tokens_priority": 1e-06, + "deprecation_date": "2027-09-02", "input_cost_per_token": 2.5e-06, "input_cost_per_token_above_272k_tokens": 5e-06, "input_cost_per_token_priority": 5e-06, @@ -6203,6 +6241,7 @@ "azure/us/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, + "deprecation_date": "2027-09-02", "input_cost_per_token": 2.75e-06, "input_cost_per_token_priority": 5.5e-06, "output_cost_per_token": 1.65e-05, @@ -6238,6 +6277,7 @@ "azure/eu/gpt-5.4-2026-03-05": { "cache_read_input_token_cost": 2.8e-07, "cache_read_input_token_cost_priority": 5.5e-07, + "deprecation_date": "2027-09-02", "input_cost_per_token": 2.75e-06, "input_cost_per_token_priority": 5.5e-06, "output_cost_per_token": 1.65e-05, @@ -6308,6 +6348,7 @@ "azure/gpt-5.4-pro-2026-03-05": { "cache_read_input_token_cost": 3e-06, "cache_read_input_token_cost_above_272k_tokens": 6e-06, + "deprecation_date": "2027-09-07", "input_cost_per_token": 3e-05, "input_cost_per_token_above_272k_tokens": 6e-05, "litellm_provider": "azure", @@ -6390,6 +6431,7 @@ "cache_read_input_token_cost_above_272k_tokens": 1e-06, "cache_read_input_token_cost_priority": 1e-06, "cache_read_input_token_cost_above_272k_tokens_priority": 2e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 5e-06, "input_cost_per_token_above_272k_tokens": 1e-05, "input_cost_per_token_priority": 1e-05, @@ -6435,6 +6477,7 @@ "cache_read_input_token_cost_above_272k_tokens": 4e-07, "cache_read_input_token_cost_priority": 4e-07, "cache_read_input_token_cost_above_272k_tokens_priority": 8e-07, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2e-06, "input_cost_per_token_above_272k_tokens": 4e-06, "input_cost_per_token_priority": 4e-06, @@ -6480,6 +6523,7 @@ "cache_read_input_token_cost_above_272k_tokens": 4e-08, "cache_read_input_token_cost_priority": 4e-08, "cache_read_input_token_cost_above_272k_tokens_priority": 8e-08, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2e-07, "input_cost_per_token_above_272k_tokens": 4e-07, "input_cost_per_token_priority": 4e-07, @@ -6566,6 +6610,7 @@ "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, @@ -6608,6 +6653,7 @@ "cache_read_input_token_cost": 2.2e-07, "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "cache_read_input_token_cost_priority": 5.5e-07, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2.2e-06, "input_cost_per_token_above_272k_tokens": 4.4e-06, "input_cost_per_token_priority": 5.5e-06, @@ -6650,6 +6696,7 @@ "cache_read_input_token_cost": 2.2e-08, "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, "input_cost_per_token_priority": 5.5e-07, @@ -6734,6 +6781,7 @@ "cache_read_input_token_cost": 5.5e-07, "cache_read_input_token_cost_above_272k_tokens": 1.1e-06, "cache_read_input_token_cost_priority": 1.375e-06, + "deprecation_date": "2028-01-11", "input_cost_per_token": 5.5e-06, "input_cost_per_token_above_272k_tokens": 1.1e-05, "input_cost_per_token_priority": 1.375e-05, @@ -6776,6 +6824,7 @@ "cache_read_input_token_cost": 2.2e-07, "cache_read_input_token_cost_above_272k_tokens": 4.4e-07, "cache_read_input_token_cost_priority": 5.5e-07, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2.2e-06, "input_cost_per_token_above_272k_tokens": 4.4e-06, "input_cost_per_token_priority": 5.5e-06, @@ -6818,6 +6867,7 @@ "cache_read_input_token_cost": 2.2e-08, "cache_read_input_token_cost_above_272k_tokens": 4.4e-08, "cache_read_input_token_cost_priority": 5.5e-08, + "deprecation_date": "2028-01-11", "input_cost_per_token": 2.2e-07, "input_cost_per_token_above_272k_tokens": 4.4e-07, "input_cost_per_token_priority": 5.5e-07, @@ -7216,6 +7266,7 @@ }, "azure/gpt-5.4-mini-2026-03-17": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2027-09-21", "input_cost_per_token": 7.5e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -7286,6 +7337,7 @@ }, "azure/gpt-5.4-nano-2026-03-17": { "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2027-09-21", "input_cost_per_token": 2e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -7321,6 +7373,7 @@ }, "azure/gpt-image-1": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-10-23", "input_cost_per_image_token": 1e-05, "input_cost_per_token": 5e-06, "litellm_provider": "azure", @@ -7432,6 +7485,7 @@ }, "azure/gpt-image-1-mini": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2027-04-07", "input_cost_per_image_token": 2.5e-06, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -7456,6 +7510,7 @@ }, "azure/gpt-image-1.5-2025-12-16": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-06-16", "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, "litellm_provider": "azure", @@ -7483,6 +7538,7 @@ }, "azure/gpt-image-2-2026-04-21": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2027-10-21", "input_cost_per_token": 5e-06, "input_cost_per_image_token": 8e-06, "litellm_provider": "azure", @@ -7613,6 +7669,7 @@ }, "azure/o1-2024-12-17": { "cache_read_input_token_cost": 7.5e-06, + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -7718,7 +7775,7 @@ "supports_vision": true }, "azure/o3-2025-04-16": { - "deprecation_date": "2026-04-16", + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -7749,6 +7806,7 @@ }, "azure/o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2026-12-26", "input_cost_per_token": 1e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -7796,6 +7854,7 @@ }, "azure/o3-mini-2025-01-31": { "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-10-01", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -7839,6 +7898,7 @@ "supports_vision": true }, "azure/o3-pro-2025-06-10": { + "deprecation_date": "2026-12-17", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -7899,6 +7959,7 @@ }, "azure/o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, + "deprecation_date": "2026-10-16", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -7939,6 +8000,7 @@ "output_cost_per_token": 0.0 }, "azure/text-embedding-3-large": { + "deprecation_date": "2028-02-09", "input_cost_per_token": 1.3e-07, "litellm_provider": "azure", "max_input_tokens": 8191, @@ -7947,7 +8009,7 @@ "output_cost_per_token": 0.0 }, "azure/text-embedding-3-small": { - "deprecation_date": "2026-04-30", + "deprecation_date": "2028-02-09", "input_cost_per_token": 2e-08, "litellm_provider": "azure", "max_input_tokens": 8191, @@ -7956,6 +8018,7 @@ "output_cost_per_token": 0.0 }, "azure/text-embedding-ada-002": { + "deprecation_date": "2028-02-09", "input_cost_per_token": 1e-07, "litellm_provider": "azure", "max_input_tokens": 8191, @@ -7987,17 +8050,19 @@ ] }, "azure/tts-1": { + "deprecation_date": "2026-12-15", "input_cost_per_character": 1.5e-05, "litellm_provider": "azure", "mode": "audio_speech" }, "azure/tts-1-hd": { + "deprecation_date": "2026-12-15", "input_cost_per_character": 3e-05, "litellm_provider": "azure", "mode": "audio_speech" }, "azure/us/gpt-4.1-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "input_cost_per_token_batches": 1.1e-06, @@ -8031,7 +8096,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-mini-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 4.4e-07, "input_cost_per_token_batches": 2.2e-07, @@ -8065,7 +8130,7 @@ "supports_web_search": false }, "azure/us/gpt-4.1-nano-2025-04-14": { - "deprecation_date": "2026-11-04", + "deprecation_date": "2026-10-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 1.1e-07, "input_cost_per_token_batches": 6e-08, @@ -8098,7 +8163,7 @@ "supports_vision": true }, "azure/us/gpt-4o-2024-08-06": { - "deprecation_date": "2026-02-27", + "deprecation_date": "2027-04-14", "cache_read_input_token_cost": 1.375e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -8115,7 +8180,7 @@ "supports_vision": true }, "azure/us/gpt-4o-2024-11-20": { - "deprecation_date": "2026-03-01", + "deprecation_date": "2027-04-14", "cache_creation_input_token_cost": 1.38e-06, "input_cost_per_token": 2.75e-06, "litellm_provider": "azure", @@ -8132,6 +8197,7 @@ }, "azure/us/gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 8.3e-08, + "deprecation_date": "2027-04-14", "input_cost_per_token": 1.65e-07, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -8213,6 +8279,7 @@ }, "azure/us/gpt-5-2025-08-07": { "cache_read_input_token_cost": 1.375e-07, + "deprecation_date": "2027-02-09", "input_cost_per_token": 1.375e-06, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -8245,6 +8312,7 @@ }, "azure/us/gpt-5-mini-2025-08-07": { "cache_read_input_token_cost": 2.75e-08, + "deprecation_date": "2027-02-09", "input_cost_per_token": 2.75e-07, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -8277,6 +8345,7 @@ }, "azure/us/gpt-5-nano-2025-08-07": { "cache_read_input_token_cost": 5.5e-09, + "deprecation_date": "2027-02-09", "input_cost_per_token": 5.5e-08, "litellm_provider": "azure", "max_input_tokens": 272000, @@ -8343,6 +8412,7 @@ }, "azure/us/gpt-5.1-chat": { "cache_read_input_token_cost": 1.4e-07, + "deprecation_date": "2026-06-29", "input_cost_per_token": 1.38e-06, "litellm_provider": "azure", "max_input_tokens": 128000, @@ -8437,6 +8507,7 @@ }, "azure/us/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, + "deprecation_date": "2026-10-21", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8481,7 +8552,7 @@ "supports_vision": false }, "azure/us/o3-2025-04-16": { - "deprecation_date": "2026-04-16", + "deprecation_date": "2026-10-21", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "azure", @@ -8512,6 +8583,7 @@ }, "azure/us/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, + "deprecation_date": "2026-10-01", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -8528,6 +8600,7 @@ }, "azure/us/o4-mini-2025-04-16": { "cache_read_input_token_cost": 3.1e-07, + "deprecation_date": "2026-10-16", "input_cost_per_token": 1.21e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8544,6 +8617,7 @@ "supports_vision": true }, "azure/whisper-1": { + "deprecation_date": "2026-12-15", "input_cost_per_second": 0.0001, "litellm_provider": "azure", "mode": "audio_transcription", @@ -11396,6 +11470,7 @@ "output_cost_per_token": 5e-07 }, "chatgpt-4o-latest": { + "deprecation_date": "2026-02-17", "input_cost_per_token": 5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -11499,6 +11574,7 @@ "cache_creation_input_token_cost": 3e-07, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-04-20", "input_cost_per_token": 2.5e-07, "litellm_provider": "anthropic", "max_input_tokens": 200000, @@ -11517,7 +11593,7 @@ "cache_creation_input_token_cost": 1.875e-05, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 1.5e-06, - "deprecation_date": "2026-05-01", + "deprecation_date": "2026-01-05", "input_cost_per_token": 1.5e-05, "litellm_provider": "anthropic", "max_input_tokens": 200000, @@ -11535,6 +11611,7 @@ "claude-4-opus-20250514": { "cache_creation_input_token_cost": 1.875e-05, "cache_read_input_token_cost": 1.5e-06, + "deprecation_date": "2026-06-15", "input_cost_per_token": 1.5e-05, "litellm_provider": "anthropic", "max_input_tokens": 200000, @@ -11563,6 +11640,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost": 3e-07, "cache_read_input_token_cost_above_200k_tokens": 6e-07, + "deprecation_date": "2026-06-15", "input_cost_per_token": 3e-06, "input_cost_per_token_above_200k_tokens": 6e-06, "litellm_provider": "anthropic", @@ -11776,7 +11854,8 @@ "supports_native_structured_output": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "deprecation_date": "2026-08-05" }, "claude-opus-4-1-20250805": { "cache_creation_input_token_cost": 1.875e-05, @@ -11812,7 +11891,7 @@ "cache_creation_input_token_cost_above_1hr": 3e-05, "cache_read_input_token_cost": 1.5e-06, "input_cost_per_token": 1.5e-05, - "deprecation_date": "2026-05-14", + "deprecation_date": "2026-06-15", "litellm_provider": "anthropic", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -12153,7 +12232,7 @@ "prompt_cache_min_tokens": 1024 }, "claude-sonnet-4-20250514": { - "deprecation_date": "2026-05-14", + "deprecation_date": "2026-06-15", "cache_creation_input_token_cost": 3.75e-06, "cache_creation_input_token_cost_above_1hr": 6e-06, "cache_read_input_token_cost": 3e-07, @@ -12508,6 +12587,7 @@ }, "codex-mini-latest": { "cache_read_input_token_cost": 3.75e-07, + "deprecation_date": "2026-02-12", "input_cost_per_token": 1.5e-06, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -12746,6 +12826,7 @@ "supports_vision": true }, "dall-e-2": { + "deprecation_date": "2026-05-12", "input_cost_per_image": 0.02, "litellm_provider": "openai", "mode": "image_generation", @@ -12756,6 +12837,7 @@ ] }, "dall-e-3": { + "deprecation_date": "2026-05-12", "input_cost_per_image": 0.04, "litellm_provider": "openai", "mode": "image_generation", @@ -15471,6 +15553,7 @@ "output_cost_per_token": 1.85e-06, "supports_function_calling": true, "supports_reasoning": true, + "supports_native_structured_output": true, "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" }, @@ -17259,6 +17342,7 @@ "supports_tool_choice": true }, "ft:babbage-002": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, "input_cost_per_token_batches": 2e-07, "litellm_provider": "text-completion-openai", @@ -17270,6 +17354,7 @@ "output_cost_per_token_batches": 2e-07 }, "ft:davinci-002": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.2e-05, "input_cost_per_token_batches": 1e-06, "litellm_provider": "text-completion-openai", @@ -17281,6 +17366,7 @@ "output_cost_per_token_batches": 1e-06 }, "ft:gpt-3.5-turbo": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "input_cost_per_token_batches": 1.5e-06, "litellm_provider": "openai", @@ -17294,6 +17380,7 @@ "supports_tool_choice": true }, "ft:gpt-3.5-turbo-0125": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-06, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -17327,6 +17414,7 @@ "supports_tool_choice": true }, "ft:gpt-4-0613": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-05, "litellm_provider": "openai", "max_input_tokens": 8192, @@ -17433,6 +17521,7 @@ }, "ft:gpt-4.1-nano-2025-04-14": { "cache_read_input_token_cost": 5e-08, + "deprecation_date": "2026-10-23", "input_cost_per_token": 2e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", @@ -17451,6 +17540,7 @@ }, "ft:o4-mini-2025-04-16": { "cache_read_input_token_cost": 1e-06, + "deprecation_date": "2026-10-23", "input_cost_per_token": 4e-06, "input_cost_per_token_batches": 2e-06, "litellm_provider": "openai", @@ -18925,6 +19015,7 @@ }, "gemini/gemini-robotics-er-1.5-preview": { "cache_read_input_token_cost": 0, + "deprecation_date": "2026-04-30", "input_cost_per_token": 3e-07, "input_cost_per_audio_token": 1e-06, "litellm_provider": "gemini", @@ -19167,6 +19258,7 @@ "uses_embed_content": true }, "gemini/gemini-embedding-001": { + "deprecation_date": "2028-05-14", "input_cost_per_token": 1.5e-07, "litellm_provider": "gemini", "max_input_tokens": 2048, @@ -19179,6 +19271,7 @@ "tpm": 10000000 }, "gemini/gemini-embedding-2-preview": { + "deprecation_date": "2026-08-10", "input_cost_per_audio_per_second": 0.00016, "input_cost_per_image": 0.00012, "input_cost_per_token": 2e-07, @@ -19388,6 +19481,7 @@ }, "gemini/gemini-2.5-flash-image": { "cache_read_input_token_cost": 3e-08, + "deprecation_date": "2026-10-02", "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", @@ -19480,6 +19574,7 @@ "supports_reasoning": false }, "gemini/gemini-3-pro-image-preview": { + "deprecation_date": "2026-06-25", "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, @@ -19565,6 +19660,7 @@ "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-flash-image-preview": { + "deprecation_date": "2026-06-25", "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "litellm_provider": "gemini", @@ -19696,6 +19792,7 @@ }, "gemini/gemini-2.5-flash-lite-preview-09-2025": { "cache_read_input_token_cost": 1e-08, + "deprecation_date": "2026-03-31", "input_cost_per_audio_token": 3e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", @@ -20077,6 +20174,7 @@ }, "gemini/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, + "deprecation_date": "2026-05-25", "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "litellm_provider": "gemini", @@ -20129,6 +20227,7 @@ "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2027-05-07", "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -20896,18 +20995,21 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "gemini/imagen-4.0-fast-generate-001": { + "deprecation_date": "2026-08-17", "litellm_provider": "gemini", "mode": "image_generation", "output_cost_per_image": 0.02, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "gemini/imagen-4.0-generate-001": { + "deprecation_date": "2026-08-17", "litellm_provider": "gemini", "mode": "image_generation", "output_cost_per_image": 0.04, "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "gemini/imagen-4.0-ultra-generate-001": { + "deprecation_date": "2026-08-17", "litellm_provider": "gemini", "mode": "image_generation", "output_cost_per_image": 0.06, @@ -20991,6 +21093,7 @@ "supports_web_search": false }, "gemini/veo-2.0-generate-001": { + "deprecation_date": "2026-06-30", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -22025,6 +22128,7 @@ "supports_vision": true }, "gpt-3.5-turbo": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 5e-07, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -22038,6 +22142,7 @@ "supports_tool_choice": true }, "gpt-3.5-turbo-0125": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 5e-07, "litellm_provider": "openai", "max_input_tokens": 16385, @@ -22097,6 +22202,7 @@ "output_cost_per_token": 2e-06 }, "gpt-4": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-05, "litellm_provider": "openai", "max_input_tokens": 8192, @@ -22137,7 +22243,7 @@ "supports_tool_choice": true }, "gpt-4-0613": { - "deprecation_date": "2025-06-06", + "deprecation_date": "2026-10-23", "input_cost_per_token": 3e-05, "litellm_provider": "openai", "max_input_tokens": 8192, @@ -22151,7 +22257,7 @@ "supports_tool_choice": true }, "gpt-4-1106-preview": { - "deprecation_date": "2026-03-26", + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22166,6 +22272,7 @@ "supports_tool_choice": true }, "gpt-4-turbo": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22182,6 +22289,7 @@ "supports_vision": true }, "gpt-4-turbo-2024-04-09": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-05, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -22363,6 +22471,7 @@ "gpt-4.1-nano": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 5e-08, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-07, "input_cost_per_token_batches": 5e-08, "input_cost_per_token_priority": 2e-07, @@ -22399,6 +22508,7 @@ "gpt-4.1-nano-2025-04-14": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 5e-08, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1e-07, "input_cost_per_token_priority": 2e-07, "input_cost_per_token_batches": 5e-08, @@ -22456,6 +22566,7 @@ "supports_vision": true }, "gpt-4o-2024-05-13": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 5e-06, "input_cost_per_token_batches": 2.5e-06, "input_cost_per_token_priority": 8.75e-06, @@ -22522,6 +22633,7 @@ "supports_vision": true }, "gpt-4o-audio-preview": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22539,6 +22651,7 @@ "supports_tool_choice": true }, "gpt-4o-audio-preview-2024-12-17": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22556,6 +22669,7 @@ "supports_tool_choice": true }, "gpt-4o-audio-preview-2025-06-03": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22573,6 +22687,7 @@ "supports_tool_choice": true }, "gpt-audio": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22642,6 +22757,7 @@ "supports_vision": false }, "gpt-audio-2025-08-28": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -22678,6 +22794,7 @@ "supports_vision": false }, "gpt-audio-mini": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -22714,6 +22831,7 @@ "supports_vision": false }, "gpt-audio-mini-2025-10-06": { + "deprecation_date": "2026-07-23", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -22837,6 +22955,7 @@ "supports_vision": true }, "gpt-4o-mini-audio-preview": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 1.5e-07, "litellm_provider": "openai", @@ -22854,6 +22973,7 @@ "supports_tool_choice": true }, "gpt-4o-mini-audio-preview-2024-12-17": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 1.5e-07, "litellm_provider": "openai", @@ -22873,6 +22993,7 @@ "gpt-4o-mini-realtime-preview": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 3e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -22892,6 +23013,7 @@ "gpt-4o-mini-realtime-preview-2024-12-17": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 3e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -22936,6 +23058,7 @@ }, "gpt-4o-mini-search-preview-2025-03-11": { "cache_read_input_token_cost": 7.5e-08, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.5e-07, "input_cost_per_token_batches": 7.5e-08, "litellm_provider": "openai", @@ -22986,6 +23109,7 @@ }, "gpt-4o-realtime-preview": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -23004,6 +23128,7 @@ }, "gpt-4o-realtime-preview-2024-12-17": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -23022,6 +23147,7 @@ }, "gpt-4o-realtime-preview-2025-06-03": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -23066,6 +23192,7 @@ }, "gpt-4o-search-preview-2025-03-11": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-07-23", "input_cost_per_token": 2.5e-06, "input_cost_per_token_batches": 1.25e-06, "litellm_provider": "openai", @@ -23098,6 +23225,7 @@ }, "gpt-image-1.5": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-12-01", "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", @@ -23112,6 +23240,7 @@ }, "gpt-image-1.5-2025-12-16": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-12-01", "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", @@ -23607,6 +23736,7 @@ "gpt-5.1-chat-latest": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -23726,6 +23856,7 @@ "gpt-5.2-chat-latest": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-08-10", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -23764,6 +23895,7 @@ "gpt-5.3-chat-latest": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-08-10", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -24647,8 +24779,8 @@ "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 400000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 272000, + "max_tokens": 272000, "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, @@ -24679,12 +24811,13 @@ "supports_minimal_reasoning_effort": true }, "gpt-5-pro-2025-10-06": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.5e-05, "input_cost_per_token_batches": 7.5e-06, "litellm_provider": "openai", "max_input_tokens": 400000, - "max_output_tokens": 128000, - "max_tokens": 128000, + "max_output_tokens": 272000, + "max_tokens": 272000, "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, @@ -24718,6 +24851,7 @@ "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 1.25e-06, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.5e-06, @@ -24793,6 +24927,7 @@ }, "gpt-5-chat-latest": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 128000, @@ -24828,6 +24963,7 @@ }, "gpt-5-codex": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -24863,6 +24999,7 @@ "gpt-5.1-codex": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_priority": 2.5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "input_cost_per_token_priority": 2.5e-06, "litellm_provider": "openai", @@ -24899,6 +25036,7 @@ }, "gpt-5.1-codex-max": { "cache_read_input_token_cost": 1.25e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", "max_input_tokens": 272000, @@ -24934,6 +25072,7 @@ "gpt-5.1-codex-mini": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-07-23", "input_cost_per_token": 2.5e-07, "input_cost_per_token_priority": 4.5e-07, "litellm_provider": "openai", @@ -24971,6 +25110,7 @@ "gpt-5.2-codex": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1.75e-06, "input_cost_per_token_priority": 3.5e-06, "litellm_provider": "openai", @@ -25088,6 +25228,7 @@ "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2.5e-07, "input_cost_per_token_flex": 1.25e-07, "input_cost_per_token_priority": 4.5e-07, @@ -25169,6 +25310,7 @@ "gpt-5-nano-2025-08-07": { "cache_read_input_token_cost": 5e-09, "cache_read_input_token_cost_flex": 2.5e-09, + "deprecation_date": "2026-12-11", "input_cost_per_token": 5e-08, "input_cost_per_token_priority": 2.5e-06, "input_cost_per_token_flex": 2.5e-08, @@ -25208,6 +25350,7 @@ }, "gpt-image-1": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-10-23", "input_cost_per_image_token": 1e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -25220,6 +25363,7 @@ }, "gpt-image-1-mini": { "cache_read_input_token_cost": 2e-07, + "deprecation_date": "2026-12-01", "input_cost_per_image_token": 2.5e-06, "input_cost_per_token": 2e-06, "litellm_provider": "openai", @@ -25233,6 +25377,7 @@ "gpt-realtime": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, @@ -25399,6 +25544,7 @@ "gpt-realtime-mini": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -25430,6 +25576,7 @@ "gpt-realtime-2025-08-28": { "cache_creation_input_audio_token_cost": 4e-07, "cache_read_input_token_cost": 4e-07, + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 3.2e-05, "input_cost_per_image": 5e-06, "input_cost_per_token": 4e-06, @@ -29258,6 +29405,7 @@ }, "o1": { "cache_read_input_token_cost": 7.5e-06, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -29277,6 +29425,7 @@ }, "o1-2024-12-17": { "cache_read_input_token_cost": 7.5e-06, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.5e-05, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -29295,6 +29444,7 @@ "supports_vision": true }, "o1-pro": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 0.00015, "input_cost_per_token_batches": 7.5e-05, "litellm_provider": "openai", @@ -29327,6 +29477,7 @@ "supports_vision": true }, "o1-pro-2025-03-19": { + "deprecation_date": "2026-10-23", "input_cost_per_token": 0.00015, "input_cost_per_token_batches": 7.5e-05, "litellm_provider": "openai", @@ -29400,6 +29551,7 @@ "cache_read_input_token_cost": 5e-07, "cache_read_input_token_cost_flex": 2.5e-07, "cache_read_input_token_cost_priority": 8.75e-07, + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-06, "input_cost_per_token_flex": 1e-06, "input_cost_per_token_priority": 3.5e-06, @@ -29436,6 +29588,7 @@ }, "o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, "litellm_provider": "openai", @@ -29470,6 +29623,7 @@ }, "o3-deep-research-2025-06-26": { "cache_read_input_token_cost": 2.5e-06, + "deprecation_date": "2026-07-23", "input_cost_per_token": 1e-05, "input_cost_per_token_batches": 5e-06, "litellm_provider": "openai", @@ -29504,6 +29658,7 @@ }, "o3-mini": { "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -29521,6 +29676,7 @@ }, "o3-mini-2025-01-31": { "cache_read_input_token_cost": 5.5e-07, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, "litellm_provider": "openai", "max_input_tokens": 200000, @@ -29568,6 +29724,7 @@ "supports_web_search": true }, "o3-pro-2025-06-10": { + "deprecation_date": "2026-12-11", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "openai", @@ -29602,6 +29759,7 @@ "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.375e-07, "cache_read_input_token_cost_priority": 5e-07, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, "input_cost_per_token_flex": 5.5e-07, "input_cost_per_token_priority": 2e-06, @@ -29627,6 +29785,7 @@ "cache_read_input_token_cost": 2.75e-07, "cache_read_input_token_cost_flex": 1.375e-07, "cache_read_input_token_cost_priority": 5e-07, + "deprecation_date": "2026-10-23", "input_cost_per_token": 1.1e-06, "input_cost_per_token_flex": 5.5e-07, "input_cost_per_token_priority": 2e-06, @@ -29650,6 +29809,7 @@ }, "o4-mini-deep-research": { "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", @@ -29684,6 +29844,7 @@ }, "o4-mini-deep-research-2025-06-26": { "cache_read_input_token_cost": 5e-07, + "deprecation_date": "2026-07-23", "input_cost_per_token": 2e-06, "input_cost_per_token_batches": 1e-06, "litellm_provider": "openai", @@ -40085,69 +40246,86 @@ }, "xai/grok-4.20-multi-agent-beta-0309": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true }, "xai/grok-4.20-beta-0309-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true }, "xai/grok-4.20-0309-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_prompt_caching": true, + "supports_response_schema": true }, "xai/grok-4.20-beta-0309-non-reasoning": { "cache_read_input_token_cost": 2e-07, - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "xai", - "max_input_tokens": 2000000, - "max_output_tokens": 2000000, - "max_tokens": 2000000, + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, "mode": "chat", - "output_cost_per_token": 6e-06, + "output_cost_per_token": 2.5e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": true + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true }, "xai/grok-4.3": { "cache_read_input_token_cost": 2e-07, @@ -40192,8 +40370,8 @@ "supports_web_search": true }, "xai/grok-4.5": { - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -40213,8 +40391,8 @@ "supports_web_search": true }, "xai/grok-4.5-latest": { - "cache_read_input_token_cost": 5e-07, - "cache_read_input_token_cost_above_200k_tokens": 1e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_above_200k_tokens": 6e-07, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, "litellm_provider": "xai", @@ -40247,51 +40425,64 @@ "supports_web_search": true }, "xai/grok-code-fast": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 2e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 2e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true, + "supports_vision": true }, "xai/grok-code-fast-1": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 2e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 2e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "deprecation_date": "2026-05-15" + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true, + "supports_vision": true }, "xai/grok-code-fast-1-0825": { - "cache_read_input_token_cost": 2e-08, - "input_cost_per_token": 2e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, "litellm_provider": "xai", "max_input_tokens": 256000, "max_output_tokens": 256000, "max_tokens": 256000, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 2e-06, "source": "https://docs.x.ai/docs/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "deprecation_date": "2026-05-15" + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true, + "supports_vision": true }, "xai/grok-vision-beta": { "input_cost_per_image": 5e-06, @@ -40331,6 +40522,7 @@ "mode": "chat", "supports_function_calling": true, "supports_reasoning": true, + "supports_native_structured_output": true, "supports_system_messages": true, "supports_tool_choice": true, "source": "https://aws.amazon.com/bedrock/pricing/" @@ -40528,6 +40720,7 @@ "mode": "chat" }, "openai/sora-2": { + "deprecation_date": "2026-09-24", "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.1, @@ -40541,6 +40734,7 @@ ] }, "openai/sora-2-pro": { + "deprecation_date": "2026-09-24", "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.3, @@ -44459,6 +44653,7 @@ ] }, "gpt-4o-mini-tts-2025-03-20": { + "deprecation_date": "2026-07-23", "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "mode": "audio_speech", @@ -44495,6 +44690,7 @@ ] }, "gpt-4o-mini-transcribe-2025-03-20": { + "deprecation_date": "2027-01-20", "input_cost_per_audio_token": 1.25e-06, "input_cost_per_token": 1.25e-06, "litellm_provider": "openai", @@ -44565,6 +44761,7 @@ "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 6e-08, + "deprecation_date": "2026-07-23", "input_cost_per_audio_token": 1e-05, "input_cost_per_image": 8e-07, "input_cost_per_token": 6e-07, @@ -44645,6 +44842,7 @@ "supports_audio_input": true }, "sora-2": { + "deprecation_date": "2026-09-24", "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.1, @@ -44658,6 +44856,7 @@ ] }, "sora-2-pro": { + "deprecation_date": "2026-09-24", "litellm_provider": "openai", "mode": "video_generation", "output_cost_per_video_per_second": 0.3, @@ -44685,6 +44884,7 @@ }, "chatgpt-image-latest": { "cache_read_input_token_cost": 1.25e-06, + "deprecation_date": "2026-12-01", "input_cost_per_image_token": 1e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -46447,6 +46647,302 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "xai/grok-4.20-0309-non-reasoning": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true + }, + "xai/grok-4.20-multi-agent-0309": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "xai", + "max_input_tokens": 1000000, + "max_output_tokens": 1000000, + "max_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_web_search": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 5e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true + }, + "xai/grok-build-0.1": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 1e-06, + "litellm_provider": "xai", + "max_input_tokens": 256000, + "max_output_tokens": 256000, + "max_tokens": 256000, + "mode": "chat", + "output_cost_per_token": 2e-06, + "source": "https://docs.x.ai/docs/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "input_cost_per_token_above_200k_tokens": 2e-06, + "output_cost_per_token_above_200k_tokens": 4e-06, + "cache_read_input_token_cost_above_200k_tokens": 4e-07, + "supports_response_schema": true, + "supports_vision": true + }, + "gpt-transcribe": { + "input_cost_per_second": 7.5e-05, + "litellm_provider": "openai", + "mode": "audio_transcription", + "source": "https://platform.openai.com/docs/models/gpt-transcribe", + "supported_endpoints": [ + "/v1/audio/transcriptions", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "gpt-live-transcribe": { + "input_cost_per_second": 0.0002833333333333333, + "litellm_provider": "openai", + "mode": "audio_transcription", + "source": "https://platform.openai.com/docs/models/gpt-live-transcribe", + "supported_endpoints": [ + "/v1/realtime", + "/v1/realtime/transcription_sessions" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "gpt-realtime-translate": { + "input_cost_per_second": 0.0005666666666666667, + "litellm_provider": "openai", + "max_input_tokens": 16000, + "max_output_tokens": 2000, + "max_tokens": 2000, + "mode": "realtime", + "source": "https://platform.openai.com/docs/models/gpt-realtime-translate", + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true + }, + "claude-mythos-5": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "prompt_cache_min_tokens": 512, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://docs.claude.com/en/docs/about-claude/models/overview", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "claude-mythos-preview": { + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 1e-05, + "litellm_provider": "anthropic", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "prompt_cache_min_tokens": 512, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "source": "https://docs.claude.com/en/docs/about-claude/models/overview", + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_output_config": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true + }, + "gemini/gemini-robotics-er-2-streaming-preview": { + "input_cost_per_audio_token": 2e-06, + "input_cost_per_token": 2e-06, + "litellm_provider": "gemini", + "mode": "chat", + "output_cost_per_token": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.014, + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014 + }, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": true, + "web_search_billing_unit": "per_query" + }, + "mistral/mistral-small-2603": { + "input_cost_per_token": 1.5e-07, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.mistral.ai/models/model-cards/mistral-small-4-0-26-03", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/labs-leanstral-1-5": { + "input_cost_per_token": 0.0, + "litellm_provider": "mistral", + "max_input_tokens": 262144, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 0.0, + "source": "https://docs.mistral.ai/models/model-cards/leanstral-1-5", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "mistral/mistral-moderation-2603": { + "input_cost_per_token": 0.0, + "litellm_provider": "mistral", + "max_input_tokens": 131072, + "mode": "moderation", + "output_cost_per_token": 0.0, + "source": "https://docs.mistral.ai/models/model-cards/mistral-moderation-26-03" + }, + "mistral/voxtral-mini-2602": { + "input_cost_per_second": 5e-05, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-transcribe-realtime-2602": { + "input_cost_per_second": 0.0001, + "litellm_provider": "mistral", + "mode": "audio_transcription", + "source": "https://docs.mistral.ai/models/model-cards/voxtral-mini-transcribe-realtime-26-02", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, + "mistral/voxtral-mini-tts-2603": { + "litellm_provider": "mistral", + "mode": "audio_speech", + "output_cost_per_character": 1.6e-05, + "source": "https://docs.mistral.ai/models/model-cards/voxtral-tts-26-03", + "supported_endpoints": [ + "/v1/audio/speech" + ], + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "audio" + ], + "supports_audio_output": true + }, "fallback_generalizations": { "rules": [ { diff --git a/osv-scanner.toml b/osv-scanner.toml index 4ef612e3a70..7ab450945f5 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -1,13 +1,3 @@ -[[IgnoredVulns]] -id = "GHSA-fwg2-594c-jp42" -ignoreUntil = 2026-08-12 -reason = "pypdf 6.15.0 (the fix) published 2026-08-06 and is still inside the P3D exclude-newer window, so uv cannot lock it yet; bump and drop this entry from 2026-08-09" - -[[IgnoredVulns]] -id = "GHSA-fp3f-mc75-235c" -ignoreUntil = 2026-08-12 -reason = "second pypdf advisory with the same 6.15.0 fix, published 2026-08-07 after the first; drop alongside GHSA-fwg2-594c-jp42 in the same bump" - [[IgnoredVulns]] id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-09-09 diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fdc81fac196..da0f608fdb5 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,18 +1,18 @@ { "ANN001": { - "limit": 3114 + "limit": 3106 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 834 + "limit": 832 }, "ANN201": { - "limit": 2031 + "limit": 2023 }, "ANN202": { - "limit": 865 + "limit": 860 }, "ANN204": { "limit": 713 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1555 + "limit": 1495 }, "ASYNC230": { "limit": 11 @@ -39,7 +39,7 @@ "limit": 505 }, "B009": { - "limit": 81 + "limit": 79 }, "B010": { "limit": 190 @@ -78,7 +78,7 @@ "limit": 1 }, "C901": { - "limit": 314 + "limit": 313 }, "D419": { "limit": 6 @@ -234,16 +234,16 @@ "limit": 5 }, "TID251": { - "limit": 1238 + "limit": 1226 }, "TRY002": { - "limit": 528 + "limit": 524 }, "TRY004": { "limit": 96 }, "TRY201": { - "limit": 407 + "limit": 405 }, "TRY203": { "limit": 113 diff --git a/ruff-strict.toml b/ruff-strict.toml index 974c49c787b..7afc5da71ee 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -16,6 +16,17 @@ external = [ "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] +[lint.per-file-ignores] +# ANN401 (explicit `Any` disallowed) has no per-line/function-level ignore mechanism +# in ruff, only file-level. These two files each have a handful of parameters that +# are genuinely heterogeneous with no fitting concrete type: a response object that +# varies across every LLM call type (completion/embedding/transcription/etc. each +# return a different shape), and *args/**kwargs forwarded verbatim with no fixed +# shape. Tried the closest existing union (CostResponseTypes) first; basedpyright +# caught a real mismatch, confirming Any is correct here, not a shortcut. +"litellm/litellm_core_utils/litellm_logging.py" = ["ANN401"] +"litellm/utils.py" = ["ANN401"] + [lint.mccabe] max-complexity = 15 diff --git a/schema.prisma b/schema.prisma index 9c871b65f40..33fd9389b63 100644 --- a/schema.prisma +++ b/schema.prisma @@ -30,7 +30,7 @@ model LiteLLM_BudgetTable { end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget tags LiteLLM_TagTable[] // multiple tags can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team - organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization } // Models on proxy @@ -893,6 +893,7 @@ model LiteLLM_DailyTeamSpend { api_requests BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + ptu_flat_cost Float @default(0.0) created_at DateTime @default(now()) updated_at DateTime @updatedAt @@ -985,6 +986,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + api_key String? + request_tags Json? @default("[]") updated_at DateTime @updatedAt updated_by String? diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index 65d0424fb5a..92eb7ef55a3 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -252,26 +252,40 @@ def scan_comments(path: Path, source: str) -> tuple[Comments, tuple[Violation, . # --------------------------------------------------------------------------- # -def mutable_names_in(annotation: ast.expr) -> Iterator[str]: +def _is_literal_subscript(node: ast.AST) -> bool: + if not isinstance(node, ast.Subscript): + return False + base: Final = node.value + return (isinstance(base, ast.Name) and base.id == "Literal") or ( + isinstance(base, ast.Attribute) and base.attr == "Literal" + ) + + +def mutable_names_in(annotation: ast.AST) -> Iterator[str]: """Yield mutable-collection names anywhere inside an annotation expression. Matches bare names (`dict`, `MutableMapping`) and dotted access (`typing.Dict`, `collections.deque`, `collections.abc.MutableMapping`), descends through nesting (`Mapping[str, list[int]]`, `tuple[set[int], ...]`) and string forward references. + Skips `Literal[...]` subtrees: their string arguments are values, not forward + references, so `Literal["list"]` is not the `list` type. """ - for node in ast.walk(annotation): - if isinstance(node, ast.Name) and node.id in MUTABLE_COLLECTIONS: - yield node.id - elif isinstance(node, ast.Attribute) and node.attr in MUTABLE_COLLECTIONS: - yield node.attr - elif isinstance(node, ast.Constant): - value: object = node.value # forward references arrive as string constants - if isinstance(value, str): - try: - inner = ast.parse(value, mode="eval").body - except SyntaxError: - continue - yield from mutable_names_in(inner) + if _is_literal_subscript(annotation): + return + if isinstance(annotation, ast.Name) and annotation.id in MUTABLE_COLLECTIONS: + yield annotation.id + elif isinstance(annotation, ast.Attribute) and annotation.attr in MUTABLE_COLLECTIONS: + yield annotation.attr + elif isinstance(annotation, ast.Constant): + value: object = annotation.value # forward references arrive as string constants + if isinstance(value, str): + try: + inner = ast.parse(value, mode="eval").body + except SyntaxError: + return + yield from mutable_names_in(inner) + for child in ast.iter_child_nodes(annotation): + yield from mutable_names_in(child) def _mutable_ann(path: Path, line: int, name: str, where: str) -> Violation: diff --git a/terraform/provider/RELEASING.md b/terraform/provider/RELEASING.md index 1dc296f29b8..7b359047e2f 100644 --- a/terraform/provider/RELEASING.md +++ b/terraform/provider/RELEASING.md @@ -106,19 +106,23 @@ Before creating a release: 4. **Land the changes in BerriAI/litellm** - Open a PR to `BerriAI/litellm` updating `terraform/provider/CHANGELOG.md` (and any source changes) and merge it. Note the merge commit SHA; the release workflow takes it as `git_ref` + Open a PR to `BerriAI/litellm` updating `terraform/provider/CHANGELOG.md` (and any source changes) and merge it ### 2. Mirror and Tag via project-releaser The provider source lives at `terraform/provider/` in `BerriAI/litellm`; `BerriAI/terraform-provider-litellm` is a thin release mirror. Do not commit or tag the mirror directly +Normally there is nothing to do here. `BerriAI/project-releaser`'s release pipeline runs the same check on every release except `adhoc`, nightly included: it reads the topmost released heading in `terraform/provider/CHANGELOG.md`, probes the mirror for `v`, and dispatches `Publish Terraform provider` only when the changelog has moved ahead of what the mirror carries. Cutting the version heading in step 1 is therefore what releases the provider, and the next release picks it up, so the wait is a day rather than a week + +Dispatch by hand only for an out-of-band release, or to recover a run that failed: + 1. Go to `BerriAI/project-releaser` > **Actions** > `Publish Terraform provider` 2. Click **Run workflow**: - `git_ref`: full 40-char commit SHA from `BerriAI/litellm` to release from - `provider_version`: the new version without the `v` prefix (e.g. `0.3.0`) - `dry_run`: optional; validates without pushing -3. The workflow rsyncs `terraform/provider/` into the mirror repo, commits, and pushes tag `v` -4. The tag push triggers the mirror's `Release` workflow (goreleaser), which is gated by the `production-release` environment approval + +Automatic or manual, the run waits on the `production-release` approval in `project-releaser`, then rsyncs `terraform/provider/` into the mirror repo, commits, and pushes tag `v`. That approval is the only one in the flow. The tag push triggers the mirror's `Release` workflow (goreleaser), which runs unattended **Important**: - Tags must follow the format: `v..` (e.g., `v0.1.2`, `v1.0.0`) diff --git a/tests/code_coverage_tests/check_prisma_binary_cache.py b/tests/code_coverage_tests/check_prisma_binary_cache.py new file mode 100644 index 00000000000..501688385ff --- /dev/null +++ b/tests/code_coverage_tests/check_prisma_binary_cache.py @@ -0,0 +1,143 @@ +"""Guard the CI cache for Prisma's CLI and engine binaries. + +``prisma generate`` shells out to ``npm install prisma@`` whenever the +prisma-client-py binary cache directory has no CLI entrypoint, pulling ~85 MB of +engines over the network. The download is normally seconds and occasionally +minutes, and a job timeout cannot tell the difference from a hung test, so an +uncached job is one slow npm response away from cancelling a passing test run. + +Three invariants keep that download off the critical path: + +1. No workflow sets ``PRISMA_BINARY_CACHE_DIR``. The prisma-client-py default is + ``~/.cache/prisma-python/binaries//``, already + keyed by both versions and the only path the cache action restores. Pointing + it elsewhere (``runner.temp`` especially, which is wiped every job) silently + guarantees a cold download. +2. Every job that generates the client also restores the cache. +3. The cache key resolves to a real version from ``uv.lock``. The action fails + the job when it cannot, so a lock format change must break here instead. +""" + +import re +import sys +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Final + +import yaml +from pydantic import BaseModel, Field, ValidationError + +REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent +WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows" +UV_LOCK: Final = REPO_ROOT / "uv.lock" +CACHE_ACTION: Final = "./.github/actions/cache-prisma-binaries" + +# Commands that reach the prisma binary cache: a direct generate, or a script +# that runs one on the caller's behalf. +PRISMA_GENERATE_MARKERS: Final = ("prisma generate", "type_check_gate.py") + + +class PrismaBinaryCacheError(Exception): + pass + + +def resolve_prisma_version(lock_text: str) -> str | None: + """Mirror of the shell lookup in the cache action's version step.""" + match: Final = re.search( + r'^name = "prisma"\n^version = "(?P[^"]+)"$', + lock_text, + re.MULTILINE, + ) + return match.group("version") if match else None + + +class WorkflowStep(BaseModel): + """The two step fields this guard reads; every other key is ignored.""" + + run: str | None = None + uses: str | None = None + + def generates_prisma_client(self) -> bool: + return self.run is not None and any(m in self.run for m in PRISMA_GENERATE_MARKERS) + + def restores_cache(self) -> bool: + return self.uses == CACHE_ACTION + + +class WorkflowJob(BaseModel): + # Absent for jobs that delegate to a reusable workflow via a job-level `uses`. + steps: tuple[WorkflowStep, ...] = () + + +class Workflow(BaseModel): + jobs: Mapping[str, WorkflowJob] = Field(default_factory=dict) + + +def parse_workflow(text: str) -> Workflow | str: + """Validate untyped YAML at the boundary so the checks below stay typed. + + Returns the parsed workflow, or a description of why it could not be read. + """ + parsed: Final = yaml.safe_load(text) + try: + return Workflow.model_validate(parsed if isinstance(parsed, dict) else {}) + except ValidationError as exc: + return f"does not parse as a workflow: {exc.error_count()} schema error(s)" + + +def lock_errors(lock_text: str) -> Iterator[str]: + if not resolve_prisma_version(lock_text): + yield ( + "uv.lock has no resolvable `prisma` package version. The version step " + f"in {CACHE_ACTION} greps the same shape and will fail every job that " + "generates the Prisma client." + ) + + +def workflow_errors(rel: Path, text: str) -> Iterator[str]: + if "PRISMA_BINARY_CACHE_DIR" in text: + yield ( + f"{rel}: sets PRISMA_BINARY_CACHE_DIR. Leave it unset so the binaries " + f"land in the version-keyed default path the {CACHE_ACTION} action restores." + ) + + workflow: Final = parse_workflow(text) + if isinstance(workflow, str): + yield f"{rel}: {workflow}" + return + + for job_name, job in workflow.jobs.items(): + if any(s.generates_prisma_client() for s in job.steps) and not any( + s.restores_cache() for s in job.steps + ): + yield ( + f"{rel}: job `{job_name}` generates the Prisma client without a " + f"`uses: {CACHE_ACTION}` step, so it downloads ~85 MB of engines " + "on every run." + ) + + +def main() -> None: + errors: Final = ( + *lock_errors(UV_LOCK.read_text()), + *( + error + for path in sorted(WORKFLOWS_DIR.glob("*.y*ml")) + for error in workflow_errors(path.relative_to(REPO_ROOT), path.read_text()) + ), + ) + + if errors: + raise PrismaBinaryCacheError( + "Prisma binary cache invariants violated:\n - " + "\n - ".join(errors) + ) + + print("Prisma binary cache invariants hold across .github/workflows/") + + +if __name__ == "__main__": + try: + main() + except PrismaBinaryCacheError as exc: + print(f"ERROR: {exc}", file=sys.stderr) + sys.exit(1) diff --git a/tests/code_coverage_tests/check_workflow_startup_safety.py b/tests/code_coverage_tests/check_workflow_startup_safety.py new file mode 100644 index 00000000000..cf150daef4c --- /dev/null +++ b/tests/code_coverage_tests/check_workflow_startup_safety.py @@ -0,0 +1,239 @@ +"""Catch workflow mistakes that GitHub reports as nothing at all. + +A workflow whose YAML is valid but whose expressions are not fails at *startup*: +the run is marked failed, no jobs are created, and no check run is ever posted. +Nothing turns red on the PR, so an entire test suite can silently stop running +while the checks list stays green. These invariants have to be enforced here +because CI cannot enforce them on itself. + +1. No arithmetic inside ``${{ }}``. GitHub expressions support grouping, index, + dereference, ``!``, the comparisons, ``&&`` and ``||``, and nothing else. A + ``${{ a + b }}`` is a startup failure, not a value. Only ``+`` and ``*`` are + flagged: ``-`` appears in hyphenated input names like ``inputs.timeout-minutes`` + and ``/`` inside ref strings, so neither can be told apart from arithmetic by + inspection alone. +2. Callers of the reusable unit-test workflow keep the job timeout at or above + the test budget plus the setup ceilings plus the runner overhead below. + Otherwise the job deadline preempts pytest inside its own advertised budget, + which is the failure the split timeouts exist to prevent, and it shows up as + a cancelled shard whose tests were passing. A budget this check cannot resolve + is reported rather than skipped, so a mistyped input or matrix column surfaces + here instead of leaving the pair silently unchecked. +""" + +import re +import sys +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import yaml +from pydantic import BaseModel, Field, ValidationError + +REPO_ROOT: Final = Path(__file__).resolve().parent.parent.parent +WORKFLOWS_DIR: Final = REPO_ROOT / ".github" / "workflows" +BASE_WORKFLOW: Final = "./.github/workflows/_test-unit-base.yml" +BASE_WORKFLOW_PATH: Final = WORKFLOWS_DIR / "_test-unit-base.yml" + +# Runner time the job clock charges but no step owns: job init, the gaps between +# steps, and post-job cleanup. Without it a job capped at exactly test + setup +# would still preempt pytest inside its own budget. +JOB_OVERHEAD_MINUTES: Final = 5 + +EXPRESSION: Final = re.compile(r"\$\{\{(?P.*?)\}\}", re.DOTALL) +QUOTED: Final = re.compile(r"'[^']*'") +ARITHMETIC: Final = re.compile(r"[+*]") +MATRIX_REF: Final = re.compile(r"^\$\{\{\s*matrix\.(?P[\w-]+)\s*\}\}$") + + +class WorkflowStartupError(Exception): + pass + + +class ReusableCall(BaseModel): + uses: str | None = None + with_: Mapping[str, object] = Field(default_factory=dict, alias="with") + strategy: Mapping[str, object] = Field(default_factory=dict) + steps: tuple[Mapping[str, object], ...] = () + + model_config = {"populate_by_name": True} + + +class WorkflowFile(BaseModel): + jobs: Mapping[str, ReusableCall] = Field(default_factory=dict) + + +def parse_workflow(text: str) -> WorkflowFile | str: + parsed: Final = yaml.safe_load(text) + try: + return WorkflowFile.model_validate(parsed if isinstance(parsed, dict) else {}) + except ValidationError as exc: + return f"does not parse as a workflow: {exc.error_count()} schema error(s)" + + +def arithmetic_expressions(text: str) -> Iterator[str]: + for match in EXPRESSION.finditer(text): + body: Final = match.group("body") + if ARITHMETIC.search(QUOTED.sub("", body)): + yield body.strip() + + +def setup_ceiling_minutes(base_text: str) -> int: + """Sum the per-step timeouts on everything the base workflow runs before pytest.""" + base: Final = yaml.safe_load(base_text) + steps: Final = base["jobs"]["run"]["steps"] + return sum( + s["timeout-minutes"] + for s in steps + if s.get("name") != "Run tests" and isinstance(s.get("timeout-minutes"), int) + ) + + +def base_default(base_text: str, name: str) -> int: + base: Final = yaml.safe_load(base_text) + return base[True]["workflow_call"]["inputs"][name]["default"] + + +@dataclass(frozen=True, slots=True) +class Column: + """A budget the caller reads from one column of its own matrix.""" + + name: str + + +def budget_source(job: ReusableCall, key: str, fallback: int) -> int | Column | str: + """A caller passes a literal, or `${{ matrix.x }}` naming a column of its matrix. + + Anything else comes back as the reason it could not be read, since a budget + nothing can resolve has to be reported rather than passed over. + """ + value: Final = job.with_.get(key) + if value is None: + return fallback + if isinstance(value, int): + return value + + matrix_ref: Final = MATRIX_REF.match(str(value)) + if not matrix_ref: + return f"passes `{key}: {value}`, which is neither a number nor a `matrix` reference." + return Column(matrix_ref.group("key")) + + +def matrix_rows(job: ReusableCall) -> Sequence[Mapping[str, object]]: + matrix: Final = job.strategy.get("matrix", {}) + entries: Final = matrix.get("include", ()) if isinstance(matrix, dict) else () + return tuple(e for e in entries if isinstance(e, dict)) + + +def budget_pairs(job: ReusableCall, test_source: int | Column, job_source: int | Column) -> Iterator[tuple[int, int]]: + """Pair each shard's test budget with the job budget of that same shard. + + Matrix-sourced budgets resolve per `include` row, so two matrix columns are + read off the same row rather than cross-producted across rows. + """ + if isinstance(test_source, int) and isinstance(job_source, int): + yield test_source, job_source + return + + for row in matrix_rows(job): + test_budget = row.get(test_source.name) if isinstance(test_source, Column) else test_source + job_budget = row.get(job_source.name) if isinstance(job_source, Column) else job_source + if isinstance(test_budget, int) and isinstance(job_budget, int): + yield test_budget, job_budget + + +def unresolved_message(where: str, job: ReusableCall, sources: Sequence[int | Column]) -> str: + """Why no shard yielded a pair of budgets to compare. + + Naming only the columns that resolve nowhere keeps the message honest: a + column every row supplies is not what left the pair unchecked. + """ + rows: Final = matrix_rows(job) + missing: Final = tuple( + f"`matrix.{s.name}`" + for s in sources + if isinstance(s, Column) and not any(isinstance(row.get(s.name), int) for row in rows) + ) + if missing: + return ( + f"{where} reads a budget from {', '.join(missing)}, which no `include` row supplies " + "as a number, so the pair would go unchecked." + ) + return ( + f"{where} reads both budgets from its matrix, but no single `include` row supplies both " + "as numbers, so the pair would go unchecked." + ) + + +def job_errors(rel: Path, job_name: str, job: ReusableCall, ceiling: int, base_text: str) -> Iterator[str]: + where: Final = f"{rel}: job `{job_name}`" + test_source: Final = budget_source(job, "timeout-minutes", base_default(base_text, "timeout-minutes")) + job_source: Final = budget_source(job, "job-timeout-minutes", base_default(base_text, "job-timeout-minutes")) + sources: Final = (test_source, job_source) + + unreadable: Final = tuple(f"{where} {reason}" for reason in sources if isinstance(reason, str)) + if unreadable: + yield from unreadable + return + + pairs: Final = tuple(budget_pairs(job, test_source, job_source)) + if not pairs: + yield unresolved_message(where, job, sources) + return + + for test_budget, job_budget in pairs: + required = test_budget + ceiling + JOB_OVERHEAD_MINUTES + if job_budget < required: + yield ( + f"{where} gives pytest {test_budget}m but caps the job at " + f"{job_budget}m. Setup can use up to {ceiling}m plus {JOB_OVERHEAD_MINUTES}m of " + f"runner overhead, so the job deadline would preempt pytest; raise " + f"job-timeout-minutes to at least {required}." + ) + + +def timeout_contract_errors(rel: Path, workflow: WorkflowFile, ceiling: int, base_text: str) -> Iterator[str]: + for job_name, job in workflow.jobs.items(): + if job.uses == BASE_WORKFLOW: + yield from job_errors(rel, job_name, job, ceiling, base_text) + + +def workflow_errors(rel: Path, text: str, ceiling: int, base_text: str) -> Iterator[str]: + for expression in arithmetic_expressions(text): + yield ( + f"{rel}: `${{{{ {expression} }}}}` uses arithmetic, which GitHub expressions do not " + "support. The workflow will fail at startup with no jobs and no check run." + ) + + workflow: Final = parse_workflow(text) + if isinstance(workflow, str): + yield f"{rel}: {workflow}" + return + + yield from timeout_contract_errors(rel, workflow, ceiling, base_text) + + +def main() -> None: + base_text: Final = BASE_WORKFLOW_PATH.read_text() + ceiling: Final = setup_ceiling_minutes(base_text) + errors: Final = tuple( + error + for path in sorted(WORKFLOWS_DIR.glob("*.y*ml")) + for error in workflow_errors(path.relative_to(REPO_ROOT), path.read_text(), ceiling, base_text) + ) + + if errors: + raise WorkflowStartupError( + "Workflow startup invariants violated:\n - " + "\n - ".join(errors) + ) + + print(f"Workflow startup invariants hold (setup ceiling {ceiling}m)") + + +if __name__ == "__main__": + try: + main() + except WorkflowStartupError as exc: + print(f"ERROR: {exc}", file=sys.stderr) + sys.exit(1) diff --git a/tests/e2e/coverage_registry/llm_claude_code_compat.yaml b/tests/e2e/coverage_registry/llm_claude_code_compat.yaml index 6edf890f7ec..c2c17a6e764 100644 --- a/tests/e2e/coverage_registry/llm_claude_code_compat.yaml +++ b/tests/e2e/coverage_registry/llm_claude_code_compat.yaml @@ -106,5 +106,5 @@ - {id: llm.messages.anthropic.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Anthropic direct"} - {id: llm.messages.azure_foundry.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: azure_foundry, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Azure AI Foundry"} - {id: llm.messages.bedrock_converse.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_converse, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Bedrock Converse"} -- {id: llm.messages.bedrock_invoke.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Bedrock Invoke"} +- {id: llm.messages.bedrock_invoke.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: bedrock_invoke, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Claude Code's client-side WebSearch tool over Bedrock Invoke; the Anthropic-managed web_search server tool is covered by llm.messages.bedrock_invoke.web_search_server_tool.nonstream.works"} - {id: llm.messages.vertex.web_search.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: vertex, capability: web_search, streaming: nonstream, assertions: [works], source: "claude_code compat matrix", rationale: "Web search server tool over Vertex AI"} diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index e8fc8067ee0..8163866abd1 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -51,6 +51,7 @@ - {id: llm.messages.anthropic.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Extended thinking via Messages API"} - {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Flagged Claude 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (#32578/#32831/#32882)", fail_before_fix: proven} - {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Claude <= 4.7 rejects role system inside messages; unflagged models must hoist reminders into top-level system or every Claude Code session 400s (#32831)", fail_before_fix: proven} +- {id: llm.messages.bedrock_invoke.web_search_server_tool.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: web_search_server_tool, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Bedrock hosts no web_search server tool, so this only works because interception rewrites it before the upstream call and the agentic loop feeds the results back in native shape; a regression that short-circuits or forwards it instead yields raw text or AWS's 400", fail_before_fix: unproven} - {id: llm.messages.azure_foundry.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: azure_foundry, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/azure_ai/anthropic/messages_transformation.py", rationale: "Azure Foundry serves Claude on the native Anthropic contract, so flagged 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (customer RCA gap)", fail_before_fix: proven} - {id: llm.messages.azure_foundry.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: azure_foundry, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/azure_ai/anthropic/messages_transformation.py", rationale: "Azure Foundry Claude <= 4.7 rejects role system inside messages; unflagged models must hoist reminders into top-level system or every Claude Code session 400s (customer RCA gap)", fail_before_fix: proven} - {id: llm.messages.vertex.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: vertex, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py", rationale: "Vertex serves Claude on the native Anthropic contract, so flagged 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (customer RCA gap)", fail_before_fix: proven} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index d17ea0e1e5e..2d921014f0e 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -73,6 +73,7 @@ LlmCapability = Literal[ "tool_use", "vision", "web_search", + "web_search_server_tool", ] diff --git a/tests/e2e/llm_translation/test_bedrock_web_search_server_tool_e2e.py b/tests/e2e/llm_translation/test_bedrock_web_search_server_tool_e2e.py new file mode 100644 index 00000000000..7ff4b5f594a --- /dev/null +++ b/tests/e2e/llm_translation/test_bedrock_web_search_server_tool_e2e.py @@ -0,0 +1,102 @@ +"""Live e2e: the Anthropic web_search server tool over Bedrock Invoke. + +Bedrock hosts none of Anthropic's ``web_search_*`` server tools, so a +``/v1/messages`` request carrying one is rejected outright with +400 "The provided request is not valid" if it reaches AWS unchanged. What makes +it work is web-search interception: the hooks rewrite the native tool into +LiteLLM's own search tool before the upstream call, Bedrock calls that tool, the +gateway runs the search, and the agentic loop feeds the results back for the +model to synthesize. The response is then rebuilt in the native shape, so a +client's citations panel sees ``server_tool_use`` and ``web_search_tool_result`` +exactly as it would from Anthropic direct. + +This cell pins that whole path. Nothing else covers it: the ``web_search`` cells +in the Claude Code compat matrix drive the CLI's *client-side* ``WebSearch`` +tool, an ordinary custom tool the CLI executes and feeds back as a +``tool_result``, and the CLI never emits a ``web_search_20250305`` definition. + +Prerequisites beyond AWS credentials: the proxy config must switch interception +on and declare a search backend. The callback entry is load-bearing; the params +block alone does not activate it. + + litellm_settings: + callbacks: ["websearch_interception"] + websearch_interception_params: + enabled_providers: ["bedrock"] + search_tool_name: e2e-search + search_tools: + - search_tool_name: e2e-search + litellm_params: + search_provider: searxng + api_base: http://127.0.0.1:8391 +""" + +from __future__ import annotations + +import pytest + +from e2e_config import unique_marker +from e2e_http import unwrap +from endpoints_client import EndpointsClient +from lifecycle import ResourceManager +from models import ( + AnthropicMessagesBody, + AnthropicWebSearchTool, + ChatMessage, + LiteLLMParamsBody, +) + +pytestmark = pytest.mark.e2e + +BEDROCK_INVOKE_BACKEND = "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0" + +WEB_SEARCH_TOOL = AnthropicWebSearchTool( + type="web_search_20250305", + name="web_search", + max_uses=3, +) + +SEARCH_PROMPT = "Use web search to tell me one recent news headline about Anthropic." + + +class TestBedrockWebSearchServerTool: + @pytest.mark.covers("llm.messages.bedrock_invoke.web_search_server_tool.nonstream.works") + def test_web_search_server_tool_is_served( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + """A bedrock deployment must answer a web_search server-tool request + instead of handing the tool to AWS and returning its 400.""" + model = f"e2e-bedrock-websearch-{unique_marker()}" + model_id = endpoints_client.create_model( + model, + LiteLLMParamsBody( + model=BEDROCK_INVOKE_BACKEND, + aws_region_name="us-east-1", + ), + ) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + response = unwrap( + endpoints_client.proxy.messages( + key, + AnthropicMessagesBody( + model=model, + max_tokens=512, + tools=[WEB_SEARCH_TOOL], + messages=[ChatMessage(role="user", content=SEARCH_PROMPT)], + ), + ) + ) + + assert response.content, f"no content blocks in response: {response}" + block_types = [block.type for block in response.content] + assert "web_search_tool_result" in block_types, ( + "the answer carries no web_search_tool_result block, so the search " + "either never ran or its results were not returned in the native shape " + f"a citations panel reads. blocks={block_types}" + ) + assert "text" in block_types, ( + "the model never synthesized an answer over the search results, so the " + f"agentic loop stopped early. blocks={block_types}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index f1c0ede0e85..bef199cdd16 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -347,23 +347,35 @@ class ToolInputSchema(BaseModel): required: list[str] = [] -class AnthropicToolSearchTool(BaseModel): - """The tool_search discovery tool. `type` carries the SDK-version-pinned - suffix (e.g. ``tool_search_tool_regex_20251119``) that LiteLLM keys its - per-provider beta-header translation on; `name` is the unsuffixed - canonical name the upstream accepts.""" +class AnthropicServerTool(BaseModel): + """An Anthropic-managed tool the upstream executes itself. It carries no + `input_schema`; `type` is the SDK-version-pinned identifier LiteLLM keys its + per-provider translation on, and `name` is the unsuffixed canonical name the + upstream accepts.""" type: str name: str +class AnthropicToolSearchTool(AnthropicServerTool): + """The tool_search discovery tool, e.g. ``tool_search_tool_regex_20251119``.""" + + +class AnthropicWebSearchTool(AnthropicServerTool): + """The web_search server tool, e.g. ``web_search_20250305``. Distinct from + Claude Code's client-side ``WebSearch`` tool, which is an ordinary custom + tool the CLI executes and feeds back as a tool_result.""" + + max_uses: int | None = None + + class AnthropicCustomTool(BaseModel): name: str description: str input_schema: ToolInputSchema -type AnthropicTool = AnthropicToolSearchTool | AnthropicCustomTool +type AnthropicTool = AnthropicToolSearchTool | AnthropicWebSearchTool | AnthropicCustomTool class AnthropicMessagesBody(BaseModel): diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index e1e5cc6c532..fde1feb80e2 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -3045,3 +3045,105 @@ async def test_same_user_different_keys_can_access_batch(): assert "batch_id" in result2 # Both keys should get the same result assert result1["batch_id"] == result2["batch_id"] + + +@pytest.mark.asyncio +async def test_file_list_cursors_are_scoped_to_the_caller(): + """A non-owner must not learn other callers' file ids through the page cursors.""" + from openai.pagination import AsyncCursorPage + from openai.types import FileObject + + from litellm.proxy._types import UserAPIKeyAuth + + owner_file = FileObject( + id="file-owner-1", + bytes=100, + created_at=1, + filename="owner.jsonl", + object="file", + purpose="batch", + status="processed", + ) + upstream_page = AsyncCursorPage[FileObject].construct( + data=[owner_file], + has_more=True, + first_id=owner_file.id, + last_id=owner_file.id, + object="list", + ) + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_many.return_value = [] + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + response = await proxy_managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth( + user_id="other-user", team_id="other-team", parent_otel_span=MagicMock() + ), + response=upstream_page, + ) + + assert response.data == [] + assert response.first_id is None + assert response.last_id is None + assert response.has_more is False + + +@pytest.mark.asyncio +async def test_file_list_cursors_follow_the_owner_scoped_page(): + from openai.pagination import AsyncCursorPage + from openai.types import FileObject + + from litellm.proxy._types import UserAPIKeyAuth + + def _raw_file(file_id: str) -> FileObject: + return FileObject( + id=file_id, + bytes=100, + created_at=1, + filename=f"{file_id}.jsonl", + object="file", + purpose="batch", + status="processed", + ) + + upstream_page = AsyncCursorPage[FileObject].construct( + data=[_raw_file("file-someone-else"), _raw_file("file-mine")], + has_more=True, + first_id="file-someone-else", + last_id="file-mine", + object="list", + ) + + managed_row = MagicMock() + managed_row.unified_file_id = "litellm_proxy:mine" + managed_row.file_object = { + "id": "file-mine", + "bytes": 100, + "created_at": 1, + "filename": "mine.jsonl", + "object": "file", + "purpose": "batch", + "status": "processed", + } + prisma_client = AsyncMock() + prisma_client.db.litellm_managedfiletable.find_many.return_value = [managed_row] + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + response = await proxy_managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth( + user_id="mine-user", parent_otel_span=MagicMock() + ), + response=upstream_page, + ) + + assert [file_object.id for file_object in response.data] == ["litellm_proxy:mine"] + assert response.first_id == "litellm_proxy:mine" + assert response.last_id == "litellm_proxy:mine" + assert response.has_more is False diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 00d5380b2f4..b13b7342c25 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -44,39 +44,87 @@ def _attrify(d: dict): return _AttrDict(d) -def _wire_batcher_for_test(prisma_client): +def _wire_batcher_for_test(prisma_client, fail_commit=False): """ Wire prisma_client.db.batch_() to return a mock batcher whose .commit() is - awaitable and whose per-table .update() calls get captured. The reset job - writes key/user/team resets via prisma.db.batch_()..update — not via - prisma_client.update_data — so tests must let that batch path complete. + awaitable and whose per-table .update()/.update_many() calls get captured. + The reset job writes every reset through prisma.db.batch_() — key/user/team + rows one by one, and the budget tier's cascade as a single transaction — so + tests must let that batch path complete. - Returns the list that will accumulate {table, where, data} dicts from - each captured update call. + Only committed batches contribute to the returned list, mirroring prisma: + with fail_commit=True the transaction blows up and must persist nothing. + + Returns the list that will accumulate {table, op, where, data} dicts from + each captured write. """ batch_calls = [] def make_batcher(): + queued = [] + class _Table: def __init__(self, table_name): self._table_name = table_name def update(self, where=None, data=None): - batch_calls.append( - {"table": self._table_name, "where": where, "data": data} + queued.append( + { + "table": self._table_name, + "op": "update", + "where": where, + "data": data, + } ) + def update_many(self, where=None, data=None): + queued.append( + { + "table": self._table_name, + "op": "update_many", + "where": where, + "data": data, + } + ) + + async def commit(): + if fail_commit: + raise RuntimeError("simulated Postgres failure committing the batch") + batch_calls.extend(queued) + batcher = MagicMock() batcher.litellm_verificationtoken = _Table("key") batcher.litellm_usertable = _Table("user") batcher.litellm_teamtable = _Table("team") - batcher.commit = AsyncMock(return_value=None) + batcher.litellm_budgettable = _Table("budget") + batcher.litellm_teammembership = _Table("team_membership") + batcher.litellm_organizationtable = _Table("org") + batcher.litellm_tagtable = _Table("tag") + batcher.litellm_endusertable = _Table("enduser") + batcher.commit = commit return batcher prisma_client.db.batch_ = MagicMock(side_effect=make_batcher) return batch_calls +def _wire_cascade_reads_for_test(prisma_client): + """ + The budget tier's cascade reads the rows it is about to zero, so their + spend counters can be invalidated after the commit. Give each of those + tables an awaitable find_many so the reads resolve instead of falling into + the job's warn-and-continue path. + """ + for table in ( + "litellm_teammembership", + "litellm_verificationtoken", + "litellm_organizationtable", + "litellm_tagtable", + "litellm_endusertable", + ): + getattr(prisma_client.db, table).find_many = AsyncMock(return_value=[]) + + @pytest.mark.asyncio async def test_reset_budget_keys_partial_failure(): """ @@ -250,41 +298,18 @@ async def test_reset_budget_users_partial_failure(): @pytest.mark.asyncio -async def test_reset_budget_endusers_partial_failure(): +async def test_reset_budget_endusers_cascade_failure_is_all_or_nothing(): """ - Test that if one enduser fails to reset, the reset loop still processes the other endusers. - We simulate six endsers where the first fails and the others are updated. + A failure anywhere in the budget-tier cascade must persist nothing, so the + tier stays due and the next scheduler tick retries it. Before the fix the + job committed the new budget_reset_at first and zeroed the dependent spend + afterwards, so a failure here left the tier stamped for the next window + while every end user stayed at the cap. """ - user1 = { - "user_id": "user1", - "spend": 20.0, - "budget_id": "budget1", - } # Will trigger simulated failure - user2 = { - "user_id": "user2", - "spend": 25.0, - "budget_id": "budget1", - } # Should be updated - user3 = { - "user_id": "user3", - "spend": 30.0, - "budget_id": "budget1", - } # Should be updated - user4 = { - "user_id": "user4", - "spend": 35.0, - "budget_id": "budget1", - } # Should be updated - user5 = { - "user_id": "user5", - "spend": 40.0, - "budget_id": "budget1", - } # Should be updated - user6 = { - "user_id": "user6", - "spend": 45.0, - "budget_id": "budget1", - } # Should be updated + endusers = [ + _attrify({"user_id": f"user{i}", "spend": 20.0 + i, "budget_id": "budget1"}) + for i in range(1, 7) + ] budget1 = LiteLLM_BudgetTableFull( **{ @@ -301,23 +326,13 @@ async def test_reset_budget_endusers_partial_failure(): if table_name == "budget": return [budget1] elif table_name == "enduser": - return [user1, user2, user3, user4, user5, user6] + return endusers return [] prisma_client.get_data = AsyncMock() prisma_client.get_data.side_effect = get_data_mock - prisma_client.update_data = AsyncMock() - # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 0} - ) - # Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 0} - ) - # Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0}) + batch_calls = _wire_batcher_for_test(prisma_client, fail_commit=True) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -326,41 +341,13 @@ async def test_reset_budget_endusers_partial_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_enduser(enduser): - if enduser["user_id"] == "user1": - raise Exception("Simulated failure for user1") - enduser["spend"] = 0.0 - return enduser + await job.reset_budget_for_litellm_budget_table() + await asyncio.sleep(0.1) - async def fake_reset_team_members(budgets_to_reset): - return 1 - - with ( - patch.object( - ResetBudgetJob, - "_reset_budget_for_enduser", - side_effect=fake_reset_enduser, - ) as mock_reset_enduser, - patch.object( - ResetBudgetJob, - "reset_budget_for_litellm_team_members", - side_effect=fake_reset_team_members, - ) as mock_reset_team_members, - ): - await job.reset_budget_for_litellm_budget_table() - await asyncio.sleep(0.1) - - assert mock_reset_enduser.call_count == 6 - assert prisma_client.update_data.await_count == 2 - update_call = prisma_client.update_data.call_args - assert update_call.kwargs.get("table_name") == "enduser" - updated_users = update_call.kwargs.get("data_list", []) - assert len(updated_users) == 5 - assert updated_users[0]["user_id"] == "user2" - assert updated_users[1]["user_id"] == "user3" - assert updated_users[2]["user_id"] == "user4" - assert updated_users[3]["user_id"] == "user5" - assert updated_users[4]["user_id"] == "user6" + assert batch_calls == [], "a failed cascade must not persist any write" + assert ( + prisma_client.update_data.await_count == 0 + ), "budget_reset_at must not be advanced outside the cascade transaction" failure_hook_calls = ( proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args_list @@ -369,6 +356,66 @@ async def test_reset_budget_endusers_partial_failure(): call.kwargs.get("call_type") == "reset_budget_endusers" for call in failure_hook_calls ) + proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() + + +@pytest.mark.asyncio +async def test_reset_budget_endusers_are_zeroed_with_the_budget_window_advance(): + """ + The happy path: every end user the tier gates is zeroed and the tier's + budget_reset_at advances, all inside one transaction. + """ + endusers = [ + _attrify({"user_id": f"user{i}", "spend": 20.0 + i, "budget_id": "budget1"}) + for i in range(1, 7) + ] + + budget1 = LiteLLM_BudgetTableFull( + **{ + "budget_id": "budget1", + "max_budget": 65.0, + "budget_duration": "2d", + "created_at": datetime.now(timezone.utc) - timedelta(days=3), + } + ) + + prisma_client = MagicMock() + + async def get_data_mock(table_name, *args, **kwargs): + if table_name == "budget": + return [budget1] + elif table_name == "enduser": + return endusers + return [] + + prisma_client.get_data = AsyncMock() + prisma_client.get_data.side_effect = get_data_mock + prisma_client.update_data = AsyncMock() + batch_calls = _wire_batcher_for_test(prisma_client) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + + job = ResetBudgetJob(proxy_logging_obj, prisma_client) + + await job.reset_budget_for_litellm_budget_table() + await asyncio.sleep(0.1) + + assert prisma_client.db.batch_.call_count == 1, "the cascade must be one transaction" + + enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] + assert len(enduser_writes) == 1 + assert enduser_writes[0]["where"]["user_id"]["in"] == [f"user{i}" for i in range(1, 7)] + assert enduser_writes[0]["data"] == {"spend": 0} + + budget_writes = [c for c in batch_calls if c["table"] == "budget"] + assert len(budget_writes) == 1 + assert budget_writes[0]["where"] == {"budget_id": "budget1"} + assert budget_writes[0]["data"]["budget_reset_at"] > datetime.now(timezone.utc) + + proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_not_called() @pytest.mark.asyncio @@ -500,16 +547,8 @@ async def test_reset_budget_continues_other_categories_on_failure(): key1, key2 = _attrify(key1), _attrify(key2) user1, user2 = _attrify(user1), _attrify(user2) team1, team2 = _attrify(team1), _attrify(team2) - # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 0} - ) - # Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 0} - ) - # Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0}) + enduser1 = _attrify(enduser1) + _wire_cascade_reads_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -541,13 +580,6 @@ async def test_reset_budget_continues_other_categories_on_failure(): ).isoformat() return team - async def fake_reset_enduser(enduser): - enduser["spend"] = 0.0 - return enduser - - async def fake_reset_team_members(budgets_to_reset): - return 1 - with ( patch.object( ResetBudgetJob, "_reset_budget_for_key", side_effect=fake_reset_key @@ -558,14 +590,6 @@ async def test_reset_budget_continues_other_categories_on_failure(): patch.object( ResetBudgetJob, "_reset_budget_for_team", side_effect=fake_reset_team ) as mock_reset_team, - patch.object( - ResetBudgetJob, "_reset_budget_for_enduser", side_effect=fake_reset_enduser - ) as mock_reset_enduser, - patch.object( - ResetBudgetJob, - "reset_budget_for_litellm_team_members", - side_effect=fake_reset_team_members, - ) as mock_reset_team_members, ): # Call the overall reset_budget method. await job.reset_budget() @@ -575,29 +599,22 @@ async def test_reset_budget_continues_other_categories_on_failure(): called_tables = { call.kwargs.get("table_name") for call in prisma_client.get_data.await_args_list } - if mock_reset_team_members.call_count > 0: - called_tables.add("team_membership") - assert called_tables == { - "key", - "user", - "team", - "budget", - "enduser", - "team_membership", - } + assert called_tables == {"key", "user", "team", "budget", "enduser"} - # After the fix, keys/users/teams write via prisma.db.batch_().
.update, - # so only budget + enduser still go through update_data. - calls = prisma_client.update_data.await_args_list - update_data_tables = [c.kwargs.get("table_name") for c in calls] - assert sorted(update_data_tables) == ["budget", "enduser"] + # Every category writes through the batch path now, so update_data is unused. + prisma_client.update_data.assert_not_awaited() - # Check enduser update: enduser succeed. - enduser_call = next(c for c in calls if c.kwargs.get("table_name") == "enduser") - assert len(enduser_call.kwargs.get("data_list", [])) == 1 + # The budget tier's cascade still ran despite the failing user category. + assert len([c for c in batch_calls if c["table"] == "team_membership"]) == 1 + enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] + assert len(enduser_writes) == 1 + assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1"]}} + assert enduser_writes[0]["data"] == {"spend": 0} # Check the new batch write path: 2 keys + 1 user (user1 failed) + 2 teams. - key_writes = [c for c in batch_calls if c["table"] == "key"] + # `op` separates the per-row resets from the cascade sweep, which also + # targets the key table. + key_writes = [c for c in batch_calls if c["table"] == "key" and c["op"] == "update"] user_writes = [c for c in batch_calls if c["table"] == "user"] team_writes = [c for c in batch_calls if c["table"] == "team"] assert len(key_writes) == 2 @@ -974,12 +991,12 @@ async def test_service_logger_teams_failure(): @pytest.mark.asyncio async def test_service_logger_endusers_success(): """ - Test that when resetting endusers succeeds the service logger success hook is called with - the correct metadata and no exception is logged. + Test that when the budget-tier cascade commits, the service logger success + hook is called with the correct metadata and no exception is logged. """ endusers = [ - {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}, - {"user_id": "user2", "spend": 25.0, "budget_id": "budget1"}, + _attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}), + _attrify({"user_id": "user2", "spend": 25.0, "budget_id": "budget1"}), ] budgets = [ LiteLLM_BudgetTableFull( @@ -1002,16 +1019,8 @@ async def test_service_logger_endusers_success(): prisma_client = MagicMock() prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() - # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 0} - ) - # Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 0} - ) - # Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0}) + batch_calls = _wire_batcher_for_test(prisma_client) + _wire_cascade_reads_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -1020,31 +1029,16 @@ async def test_service_logger_endusers_success(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_enduser(enduser): - enduser["spend"] = 0.0 - return enduser + with patch( + "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception" + ) as mock_verbose_exc: + await job.reset_budget_for_litellm_budget_table() + await asyncio.sleep(0.1) + mock_verbose_exc.assert_not_called() - async def fake_reset_team_members(budgets_to_reset): - return 1 - - with ( - patch.object( - ResetBudgetJob, - "_reset_budget_for_enduser", - side_effect=fake_reset_enduser, - ) as mock_reset_enduser, - patch.object( - ResetBudgetJob, - "reset_budget_for_litellm_team_members", - side_effect=fake_reset_team_members, - ) as mock_reset_team_members, - ): - with patch( - "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception" - ) as mock_verbose_exc: - await job.reset_budget_for_litellm_budget_table() - await asyncio.sleep(0.1) - mock_verbose_exc.assert_not_called() + enduser_writes = [c for c in batch_calls if c["table"] == "enduser"] + assert len(enduser_writes) == 1 + assert enduser_writes[0]["where"] == {"user_id": {"in": ["user1", "user2"]}} proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_called_once() ( @@ -1062,12 +1056,12 @@ async def test_service_logger_endusers_success(): @pytest.mark.asyncio async def test_service_logger_endusers_failure(): """ - Test that a failure during enduser reset calls the failure hook with appropriate metadata, - logs the exception, and does not call the success hook. + Test that a failed cascade calls the failure hook with the rows it had + found, logs the exception, and does not call the success hook. """ endusers = [ - {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}, - {"user_id": "user2", "spend": 25.0, "budget_id": "budget1"}, + _attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}), + _attrify({"user_id": "user2", "spend": 25.0, "budget_id": "budget1"}), ] budgets = [ LiteLLM_BudgetTableFull( @@ -1090,16 +1084,8 @@ async def test_service_logger_endusers_failure(): prisma_client = MagicMock() prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() - # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 0} - ) - # Mock db.litellm_organizationtable.update_many (used by reset_budget_for_orgs_linked_to_budgets) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 0} - ) - # Mock db.litellm_tagtable.update_many (used by reset_budget_for_tags_linked_to_budgets) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0}) + _wire_batcher_for_test(prisma_client, fail_commit=True) + _wire_cascade_reads_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -1108,39 +1094,16 @@ async def test_service_logger_endusers_failure(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_enduser(enduser): - if enduser["user_id"] == "user1": - raise Exception("Simulated failure for user1") - enduser["spend"] = 0.0 - return enduser - - async def fake_reset_team_members(budgets_to_reset): - return 1 - - with ( - patch.object( - ResetBudgetJob, - "_reset_budget_for_enduser", - side_effect=fake_reset_enduser, - ) as mock_reset_enduser, - patch.object( - ResetBudgetJob, - "reset_budget_for_litellm_team_members", - side_effect=fake_reset_team_members, - ) as mock_reset_team_members, - ): - with patch( - "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception" - ) as mock_verbose_exc: - await job.reset_budget_for_litellm_budget_table() - await asyncio.sleep(0.1) - # Verify exception logging - assert mock_verbose_exc.call_count >= 1 - # Verify exception was logged with correct message - assert any( - "Failed to reset budget for enduser" in str(call.args) - for call in mock_verbose_exc.call_args_list - ) + with patch( + "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception" + ) as mock_verbose_exc: + await job.reset_budget_for_litellm_budget_table() + await asyncio.sleep(0.1) + # The log must name the whole cascade, not just end users: the write + # that failed could have been any of team member / enduser / org / tag + # spend or the budget_reset_at advance. + assert mock_verbose_exc.call_count == 1 + assert "budget table cascade" in str(mock_verbose_exc.call_args.args[0]) proxy_logging_obj.service_logging_obj.async_service_failure_hook.assert_called_once() ( @@ -1158,8 +1121,8 @@ async def test_service_logger_endusers_failure(): @pytest.mark.asyncio async def test_reset_budget_for_litellm_team_members_called(): """ - Test that when reset_budget_for_litellm_budget_table is called, - team members' budgets are also reset via reset_budget_for_litellm_team_members + Test that when reset_budget_for_litellm_budget_table is called, team + members' spend is zeroed as part of the cascade transaction. """ # Arrange budget1 = LiteLLM_BudgetTableFull( @@ -1171,7 +1134,7 @@ async def test_reset_budget_for_litellm_team_members_called(): } ) - enduser1 = {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"} + enduser1 = _attrify({"user_id": "user1", "spend": 25.0, "budget_id": "budget1"}) prisma_client = MagicMock() @@ -1184,20 +1147,9 @@ async def test_reset_budget_for_litellm_team_members_called(): prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() - - # Mock the db.litellm_teammembership.update_many call prisma_client.db = MagicMock() - prisma_client.db.litellm_teammembership = MagicMock() - prisma_client.db.litellm_teammembership.update_many = AsyncMock( - return_value={"count": 2} - ) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( - return_value={"count": 0} - ) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock( - return_value={"count": 0} - ) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 0}) + batch_calls = _wire_batcher_for_test(prisma_client) + _wire_cascade_reads_for_test(prisma_client) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -1206,23 +1158,11 @@ async def test_reset_budget_for_litellm_team_members_called(): job = ResetBudgetJob(proxy_logging_obj, prisma_client) - async def fake_reset_enduser(enduser): - enduser["spend"] = 0.0 - return enduser - - with patch.object( - ResetBudgetJob, - "_reset_budget_for_enduser", - side_effect=fake_reset_enduser, - ): - # Act - await job.reset_budget_for_litellm_budget_table() + # Act + await job.reset_budget_for_litellm_budget_table() # Assert - # Verify that the team membership update was called - prisma_client.db.litellm_teammembership.update_many.assert_called_once() - - # Verify the call was made with correct parameters - call_args = prisma_client.db.litellm_teammembership.update_many.call_args - assert call_args.kwargs["where"]["budget_id"]["in"] == ["budget1"] - assert call_args.kwargs["data"]["spend"] == 0 + team_member_writes = [c for c in batch_calls if c["table"] == "team_membership"] + assert len(team_member_writes) == 1 + assert team_member_writes[0]["where"]["budget_id"]["in"] == ["budget1"] + assert team_member_writes[0]["data"] == {"spend": 0} diff --git a/tests/logging_callback_tests/test_standard_logging_payload.py b/tests/logging_callback_tests/test_standard_logging_payload.py index f29b245b3be..d13cdf1337a 100644 --- a/tests/logging_callback_tests/test_standard_logging_payload.py +++ b/tests/logging_callback_tests/test_standard_logging_payload.py @@ -471,8 +471,8 @@ def test_get_final_response_obj(): litellm.turn_off_message_logging = False -def test_get_standard_logging_payload_trace_id(): - """Test _get_standard_logging_payload_trace_id with different input scenarios""" +def testget_standard_logging_payload_trace_id(): + """Test get_standard_logging_payload_trace_id with different input scenarios""" # Test case 1: When litellm_trace_id is provided in litellm_params from unittest.mock import MagicMock @@ -482,33 +482,134 @@ def test_get_standard_logging_payload_trace_id(): # Test when litellm_trace_id is in litellm_params litellm_params = {"litellm_trace_id": "dynamic-trace-id"} - result = StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( logging_obj=mock_logging_obj, litellm_params=litellm_params ) assert result == "dynamic-trace-id" # Test case 2: When litellm_trace_id is not provided in litellm_params litellm_params = {} - result = StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( logging_obj=mock_logging_obj, litellm_params=litellm_params ) assert result == "default-trace-id" # Test case 3: When litellm_params is None - result = StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( logging_obj=mock_logging_obj, litellm_params={} ) assert result == "default-trace-id" # Test case 4: When litellm_trace_id in params is not a string litellm_params = {"litellm_trace_id": 12345} - result = StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( logging_obj=mock_logging_obj, litellm_params=litellm_params ) assert result == "12345" assert isinstance(result, str) +def testget_standard_logging_payload_trace_id_prioritizes_trace_id_when_flag_on(monkeypatch): + """With request_correlation_in_logs on, an explicit litellm_trace_id wins over litellm_session_id.""" + from unittest.mock import MagicMock + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "default-trace-id" + + litellm_params = {"litellm_trace_id": "the-trace-id", "litellm_session_id": "the-session-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "the-trace-id" + + +def testget_standard_logging_payload_trace_id_prioritizes_session_id_when_flag_off(monkeypatch): + """With request_correlation_in_logs off (default), legacy behavior is preserved: + litellm_session_id still wins over litellm_trace_id.""" + from unittest.mock import MagicMock + + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "default-trace-id" + + litellm_params = {"litellm_trace_id": "the-trace-id", "litellm_session_id": "the-session-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "the-session-id" + + +def testget_standard_logging_payload_session_id_when_flag_on(monkeypatch): + """Test get_standard_logging_payload_session_id with different input scenarios, flag enabled""" + from unittest.mock import MagicMock + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_session_id = "" + + # Test case 1: litellm_session_id provided directly in litellm_params + litellm_params = {"litellm_session_id": "dynamic-session-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "dynamic-session-id" + + # Test case 2: falls back to metadata.session_id when not in litellm_params directly + litellm_params = {"metadata": {"session_id": "metadata-session-id"}} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "metadata-session-id" + + # Test case 3: falls back to logging_obj.litellm_session_id when nothing else is set + mock_logging_obj.litellm_session_id = "obj-session-id" + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params={} + ) + assert result == "obj-session-id" + + # Test case 4: empty string when no session id was supplied anywhere + mock_logging_obj.litellm_session_id = "" + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params={} + ) + assert result == "" + + # Test case 5: non-string session id in params is coerced to str + litellm_params = {"litellm_session_id": 98765} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "98765" + assert isinstance(result, str) + + # Test case 6: trace_id and session_id are independent - passing only a trace id + # must not populate session_id + litellm_params = {"litellm_trace_id": "some-trace-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "" + + +def testget_standard_logging_payload_session_id_empty_when_flag_off(monkeypatch): + """When request_correlation_in_logs is off (default), session_id is always empty, + even if litellm_session_id was explicitly supplied - preserves the pre-existing + StandardLoggingPayload shape for callers who haven't opted in.""" + from unittest.mock import MagicMock + + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_session_id = "obj-session-id" + + litellm_params = {"litellm_session_id": "dynamic-session-id"} + result = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + assert result == "" + + def test_truncate_standard_logging_payload(): """ 1. original messages, response, and error_str should NOT BE MODIFIED, since these are from kwargs diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 800c97d7ba1..6d7ada17ec5 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1641,3 +1641,153 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] +class TestBatchCostAttribution: + """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so + the batch-cost log is attributed like a non-batch request.""" + + def _instance(self, key_row=None, team_row=None, user_row=None): + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + prisma = MagicMock() + prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_row) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) + return CheckBatchCost( + proxy_logging_obj=MagicMock(), + prisma_client=prisma, + llm_router=MagicMock(), + ) + + def _job(self, **overrides): + from types import SimpleNamespace + + fields = { + "created_by": "alice", + "team_id": "team-alpha", + "api_key": "hash-alice", + "request_tags": ["env:prod"], + } + fields.update(overrides) + return SimpleNamespace(unified_object_id="uoi", **fields) + + @pytest.mark.asyncio + async def test_metadata_carries_key_team_and_tags(self): + """The spend row names the creating key, its team, both aliases, and the tags.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key"), + team_row=SimpleNamespace(team_alias="Team Alpha"), + user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key"] == "hash-alice" + assert metadata["user_api_key_user_id"] == "alice" + assert metadata["user_api_key_team_id"] == "team-alpha" + assert metadata["user_api_key_alias"] == "prod-key" + assert metadata["user_api_key_team_alias"] == "Team Alpha" + assert metadata["tags"] == ["env:prod"] + + @pytest.mark.asyncio + async def test_metadata_tolerates_legacy_row_without_columns(self): + """Rows created before the columns existed carry only created_by/team_id and must + still produce an attributed row rather than raising.""" + instance = self._instance() + job = self._job(api_key=None, request_tags=None) + + metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + + assert metadata["user_api_key"] is None + assert metadata["user_api_key_user_id"] == "alice" + assert metadata["user_api_key_team_id"] == "team-alpha" + assert "tags" not in metadata + + @pytest.mark.asyncio + async def test_metadata_keeps_key_when_team_key_has_no_user(self): + """A team-scoped key carries no user id. The user lookup is skipped (prisma rejects + a None user_id) and the key hash still drives key-level attribution.""" + from types import SimpleNamespace + + instance = self._instance(key_row=SimpleNamespace(key_alias="svc-key")) + job = self._job(created_by=None) + + metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + + assert metadata["user_api_key"] == "hash-alice" + assert metadata["user_api_key_user_id"] is None + assert metadata["user_api_key_alias"] == "svc-key" + instance.prisma_client.db.litellm_usertable.find_unique.assert_not_called() + + @pytest.mark.asyncio + async def test_metadata_drops_non_string_tags(self): + """Non-string tags are dropped so a malformed stored value cannot slip past the + tag-budget checks that consume this metadata.""" + instance = self._instance() + job = self._job(request_tags=["env:prod", 7, None, "team:ml"]) + + metadata = await instance._build_creator_attribution_metadata(job, "batch-1") + + assert metadata["tags"] == ["env:prod", "team:ml"] + + @pytest.mark.asyncio + async def test_key_alias_lookup_failure_does_not_break_attribution(self): + """An alias lookup failure must not lose the spend row; the key hash and team still + attribute it.""" + instance = self._instance() + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=Exception("db down") + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key"] == "hash-alice" + assert metadata.get("user_api_key_alias") is None + + @pytest.mark.asyncio + async def test_unnamed_key_keeps_the_creating_user_alias(self): + """Regression: a key generated without key_alias resolves to no alias, and the + overwrite must not null out the creating user's alias that _get_user_info supplied. + Most keys carry no alias, so this is the common batch, not an edge case.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias=None), + user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_alias"] == "Alice Chen" + assert metadata["user_api_key"] == "hash-alice" + + @pytest.mark.asyncio + async def test_rotated_key_keeps_the_creating_user_alias(self): + """Batches outlive keys. When the creating key has been rotated or deleted the + lookup returns no row, and the spend log keeps a resolvable name instead of null.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=None, + user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_alias"] == "Alice Chen" + + @pytest.mark.asyncio + async def test_named_key_still_owns_the_alias(self): + """The fallback must not weaken the intended precedence: a key that has its own + alias still overrides the creating user's.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key"), + user_row=SimpleNamespace(user_email="alice@example.com", user_alias="Alice Chen"), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_alias"] == "prod-key" diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index f64994cb3b1..bfbc92adc74 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2697,6 +2697,79 @@ def test_get_timeout_from_request(): assert timeout == 90.5 +def test_add_litellm_data_for_backend_llm_call_marks_client_side_timeout(): + """A caller-supplied x-litellm-timeout must be marked with client_side_timeout=True, + so the router's fallback-cooldown trigger can tell it apart from a deployment + actually timing out (a caller could otherwise force every deployment in a fallback + chain to look unhealthy with a single near-zero timeout request).""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key") + + data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={"x-litellm-timeout": "0.001"}, + request_data={}, + user_api_key_dict=user_api_key_dict, + ) + assert data["timeout"] == 0.001 + assert data["client_side_timeout"] is True + + data_without_header = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={}, + request_data={}, + user_api_key_dict=user_api_key_dict, + ) + assert "client_side_timeout" not in data_without_header + + +@pytest.mark.parametrize( + "request_data", + [ + {"timeout": 0.001}, + {"request_timeout": 0.001}, + {"stream_timeout": 0.001}, + ], +) +def test_add_litellm_data_for_backend_llm_call_marks_client_side_timeout_from_body( + request_data, +): + """Router._get_timeout resolves the effective timeout from kwargs["timeout"], + kwargs["request_timeout"], or kwargs["stream_timeout"], and a caller can supply any + of those directly in the request body, not just via the x-litellm-timeout header. + Missing this would let a caller force a 408 on every deployment in a fallback chain + without it being recognized as caller-controlled, cooling down deployments other + tenants rely on.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key") + + data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={}, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + ) + assert data["client_side_timeout"] is True + + +def test_add_litellm_data_for_backend_llm_call_ignores_forged_client_side_timeout(): + """The caller-supplied client_side_timeout key itself must never be trusted verbatim: + the marker is always recomputed from the actual timeout sources, so a caller can't + forge client_side_timeout=True to dodge cooldown on a real deployment failure.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key") + + data = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={}, + request_data={"client_side_timeout": True}, + user_api_key_dict=user_api_key_dict, + ) + assert "client_side_timeout" not in data + + @pytest.mark.parametrize( "ui_exists, ui_has_content", [ diff --git a/tests/proxy_unit_tests/test_update_daily_tag_spend.py b/tests/proxy_unit_tests/test_update_daily_tag_spend.py index 80616ade5ef..35e9c6796eb 100644 --- a/tests/proxy_unit_tests/test_update_daily_tag_spend.py +++ b/tests/proxy_unit_tests/test_update_daily_tag_spend.py @@ -91,17 +91,13 @@ async def test_daily_tag_spend_retries_then_succeeds(): prisma_client = MagicMock() proxy_logging_obj = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_batcher.litellm_dailytagspend = mock_table - - # Fail entering batch context 3 times with retryable DB errors, then succeed. - prisma_client.db.batch_.return_value.__aenter__ = AsyncMock( + # Fail the upsert 3 times with retryable DB errors, then succeed. + prisma_client.db.execute_raw = AsyncMock( side_effect=[ httpx.ConnectError("x"), httpx.ConnectError("x"), httpx.ConnectError("x"), - mock_batcher, + 1, ] ) @@ -138,6 +134,10 @@ async def test_daily_tag_spend_retries_then_succeeds(): daily_spend_transactions=daily_spend_transactions, ) - assert prisma_client.db.batch_.return_value.__aenter__.await_count == 4 + assert prisma_client.db.execute_raw.await_count == 4 assert sleep_mock.await_count == 3 - mock_table.upsert.assert_called_once() + # The batch is one statement, so the successful attempt is a single call carrying + # the row rather than one call per key. + final_sql = prisma_client.db.execute_raw.await_args.args[0] + assert final_sql.count("ON CONFLICT") == 1 + assert "prod-tag" in prisma_client.db.execute_raw.await_args.args[1:] diff --git a/tests/router_unit_tests/test_router_cooldown_per_deployment.py b/tests/router_unit_tests/test_router_cooldown_per_deployment.py new file mode 100644 index 00000000000..b8ae8a8c013 --- /dev/null +++ b/tests/router_unit_tests/test_router_cooldown_per_deployment.py @@ -0,0 +1,779 @@ +""" +Tests for per-deployment cooldown policy overrides, DualCache TTL correction, +and fallback-path cooldown gap fix. +""" + +import time +from unittest.mock import MagicMock, patch + +import pytest + +import litellm +from litellm import Router +from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue +from litellm.router_utils.cooldown_handlers import ( + _get_deployment_cooldown_policy, + _has_explicit_allowed_fails_policy_for_exception, + _resolve_allowed_fails_from_policy, + _should_cooldown_deployment, + mark_advisor_orchestration_failure, + should_cooldown_based_on_allowed_fails_policy, +) +from litellm.router_utils.fallback_event_handlers import _trigger_cooldown_for_failed_deployment +from litellm.types.router import AllowedFailsPolicy + + +def _make_router(model_list: list, **kwargs) -> Router: + return Router(model_list=model_list, **kwargs) + + +class TestDeploymentLevelAllowedFails: + def test_deployment_level_allowed_fails_overrides_router_level(self): + """ + A deployment with model_info.allowed_fails=0 must enter cooldown after 1 + failure even when the router-level allowed_fails=10. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": { + "id": "primary", + "allowed_fails": 0, + }, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": {"id": "secondary"}, + }, + ], + allowed_fails=10, + ) + + _exception = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + should_cooldown = _should_cooldown_deployment( + litellm_router_instance=router, + deployment="primary", + exception_status=429, + original_exception=_exception, + ) + + assert should_cooldown is True, "Deployment-level allowed_fails=0 should force cooldown after first failure" + + def test_deployment_level_allowed_fails_does_not_affect_other_deployments(self): + """ + A deployment without model_info.allowed_fails must still use the router-level + allowed_fails and not be pulled into cooldown prematurely. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": { + "id": "primary", + "allowed_fails": 0, + }, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": {"id": "secondary"}, + }, + ], + allowed_fails=10, + ) + + _exception = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + should_cooldown = _should_cooldown_deployment( + litellm_router_instance=router, + deployment="secondary", + exception_status=429, + original_exception=_exception, + ) + + assert should_cooldown is False, ( + "secondary has no deployment-level policy; with allowed_fails=10 it should not cool down on first failure" + ) + + +class TestDeploymentLevelAllowedFailsPolicyByExceptionType: + def test_rate_limit_error_triggers_cooldown_with_zero_threshold(self): + """ + RateLimitErrorAllowedFails=0 must trigger cooldown after 1 RateLimitError + even when allowed_fails=5 for other exception types. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": { + "id": "primary", + "allowed_fails_policy": { + "RateLimitErrorAllowedFails": 0, + "InternalServerErrorAllowedFails": 5, + }, + }, + }, + ], + allowed_fails=10, + ) + + rate_limit_exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + should_cooldown = _should_cooldown_deployment( + litellm_router_instance=router, + deployment="primary", + exception_status=429, + original_exception=rate_limit_exc, + ) + + assert should_cooldown is True, "RateLimitErrorAllowedFails=0 must trigger cooldown on first rate limit error" + + def test_internal_server_error_respects_per_exception_threshold(self): + """ + InternalServerErrorAllowedFails=5 must allow 5 InternalServerErrors before cooldown. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": { + "id": "primary", + "allowed_fails_policy": { + "RateLimitErrorAllowedFails": 0, + "InternalServerErrorAllowedFails": 5, + }, + }, + }, + ], + allowed_fails=10, + ) + + ise = litellm.InternalServerError("Internal error", "openai", "gpt-4") + + for _ in range(5): + should_cooldown = _should_cooldown_deployment( + litellm_router_instance=router, + deployment="primary", + exception_status=500, + original_exception=ise, + ) + assert should_cooldown is False, "Should not cooldown within the allowed_fails threshold" + + should_cooldown = _should_cooldown_deployment( + litellm_router_instance=router, + deployment="primary", + exception_status=500, + original_exception=ise, + ) + assert should_cooldown is True, "Should cooldown after exceeding InternalServerErrorAllowedFails=5" + + +class TestExceptionTypeCountersTrackedIndependently: + def test_cache_key_suffix_separates_exception_type_counters(self): + """ + When cache_key_suffix is provided, fail counters for different exception types + must be independent; RateLimitError fails must not bleed into generic counters. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": {"id": "primary"}, + }, + ], + allowed_fails=10, + ) + + rate_limit_exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + ise = litellm.InternalServerError("Internal error", "openai", "gpt-4") + + for _ in range(3): + should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="primary", + original_exception=rate_limit_exc, + allowed_fails_override=5, + 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 + + assert rl_counter == 3, "RateLimitError counter should be 3" + assert generic_counter == 0, "generic counter must be untouched by RateLimitError increments" + + should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="primary", + original_exception=ise, + allowed_fails_override=5, + 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 + + assert generic_counter_after == 1, "generic counter should now be 1" + assert rl_counter_after == 3, "RateLimitError counter must remain unchanged after InternalServerError" + + +class TestCooldownCacheTTLCorrection: + def _make_cooldown_cache(self) -> CooldownCache: + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory) + return CooldownCache(cache=dual_cache, default_cooldown_time=60.0) + + def test_expired_entry_evicted_and_not_returned(self): + """ + An entry with timestamp+cooldown_time in the past must be evicted from + in-memory cache and excluded from the active cooldown list. + """ + cc = self._make_cooldown_cache() + model_id = "expired-deployment" + key = CooldownCache.get_cooldown_cache_key(model_id) + + expired_value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time() - 120.0, + "cooldown_time": 60.0, + } + cc.cache.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" + + def test_active_entry_is_returned(self): + """ + An entry whose cooldown window has not elapsed must appear in the active list. + """ + cc = self._make_cooldown_cache() + model_id = "active-deployment" + key = CooldownCache.get_cooldown_cache_key(model_id) + + active_value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time(), + "cooldown_time": 60.0, + } + cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + + active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) + + assert len(active) == 1 + assert active[0][0] == model_id + + def test_ttl_corrected_when_in_memory_expiry_far_exceeds_remaining(self): + """ + When DualCache backfills from Redis using the default 600s TTL, the in-memory + TTL must be corrected to min(remaining, 60) seconds. + """ + cc = self._make_cooldown_cache() + model_id = "backfilled-deployment" + key = CooldownCache.get_cooldown_cache_key(model_id) + + remaining = 30.0 + value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time() - (60.0 - remaining), + "cooldown_time": 60.0, + } + cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + + before_expiry = cc.cache.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) + assert after_expiry is not None + corrected_remaining = after_expiry - time.time() + assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" + assert corrected_remaining > 0, "Corrected TTL must be positive (cooldown still active)" + + @pytest.mark.asyncio + async def test_async_expired_entry_evicted(self): + """ + Async path must also evict expired entries. + """ + cc = self._make_cooldown_cache() + model_id = "async-expired" + key = CooldownCache.get_cooldown_cache_key(model_id) + + expired_value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time() - 120.0, + "cooldown_time": 60.0, + } + cc.cache.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 + + +class TestFallbackDeploymentCooldown: + def test_trigger_cooldown_for_failed_deployment_calls_set_cooldown(self): + """ + _trigger_cooldown_for_failed_deployment must call _set_cooldown_deployments + with the deployment ID stamped on the exception. + """ + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs={}, + exception=exc, + ) + + mock_set_cooldown.assert_called_once() + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["deployment"] == "fallback-deployment" + assert call_kwargs["original_exception"] is exc + + def test_trigger_cooldown_no_op_when_deployment_id_missing(self): + """ + _trigger_cooldown_for_failed_deployment must not raise and must skip + _set_cooldown_deployments when the exception has no failed_deployment_id. + """ + mock_router = MagicMock() + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs={}, + exception=RuntimeError("no stamped deployment id"), + ) + + mock_set_cooldown.assert_not_called() + + def test_trigger_cooldown_does_not_trust_caller_supplied_metadata_bucket(self): + """ + A metadata bucket can't reliably be told apart from a caller-supplied one + without knowing the call's function_name, so a client with permission to + set metadata must not be able to get an arbitrary deployment cooled down + by forging a deployment_model_name marker. + """ + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + kwargs = { + "metadata": { + "model_info": {"id": "attacker-chosen-deployment"}, + "deployment_model_name": "gpt-4", + } + } + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs=kwargs, + exception=exc, + ) + + mock_set_cooldown.assert_not_called() + + def test_trigger_cooldown_increments_failure_counter_before_cooldown_check(self): + """ + The fallback path must feed the same per-minute failure counter the + primary path uses, or repeated fallback failures never accumulate toward + the default percent-fail-rate cooldown threshold. + """ + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with ( + patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown, + patch( + "litellm.router_utils.fallback_event_handlers.increment_deployment_failures_for_current_minute" + ) as mock_increment, + ): + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + mock_increment.assert_called_once_with( + litellm_router_instance=mock_router, deployment_id="fallback-deployment" + ) + mock_set_cooldown.assert_called_once() + + def test_trigger_cooldown_uses_deployment_cooldown_time_override(self): + """ + When the deployment has a model_info.cooldown_time, that value must be + passed as time_to_cooldown rather than the router-level cooldown_time. + """ + mock_router = MagicMock() + mock_router.cooldown_time = 300.0 + mock_router.get_model_info.return_value = {"model_info": {"cooldown_time": 30.0}} + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs={}, + exception=exc, + ) + + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["time_to_cooldown"] == 30.0, ( + "Deployment-level cooldown_time must override router-level value" + ) + + def test_trigger_cooldown_skipped_for_advisor_orchestration_failure(self): + """ + A failure tagged as originating from advisor orchestration (not the selected + deployment) must not cool down the fallback deployment, matching the same + guard already applied in Router.deployment_callback_on_failure. + """ + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + mark_advisor_orchestration_failure(exc) + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs={}, + exception=exc, + ) + + mock_set_cooldown.assert_not_called() + + def test_trigger_cooldown_falls_back_to_litellm_params_cooldown_time(self): + """ + cooldown_time has pre-existing litellm_params support on the primary + failure path (Router.deployment_callback_on_failure), so it must still be + honored as a fallback when model_info doesn't set it, unlike the new + allowed_fails/allowed_fails_policy fields which are model_info-only. + """ + mock_router = MagicMock() + mock_router.cooldown_time = 300.0 + mock_router.get_model_info.return_value = {"litellm_params": {"cooldown_time": 30.0}} + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs={}, + exception=exc, + ) + + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["time_to_cooldown"] == 30.0, ( + "litellm_params.cooldown_time must still be honored as a fallback" + ) + + def test_trigger_cooldown_prefers_model_info_cooldown_time_over_litellm_params(self): + mock_router = MagicMock() + mock_router.cooldown_time = 300.0 + mock_router.get_model_info.return_value = { + "model_info": {"cooldown_time": 15.0}, + "litellm_params": {"cooldown_time": 30.0}, + } + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs={}, + exception=exc, + ) + + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["time_to_cooldown"] == 15.0, "model_info.cooldown_time must take priority" + + +class TestSingleDeploymentModelGroupProtection: + def test_generic_allowed_fails_does_not_bypass_single_deployment_protection(self): + """ + Setting only a generic model_info.allowed_fails on a single-deployment model + group must not disable the "avoid cooldowns on single deployment model groups" + safety net; before this feature existed the field had no effect at all here, + so a plain 500 error must behave the same as the no-policy control. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": {"id": "solo", "allowed_fails": 1}, + }, + ], + ) + + exc = Exception("Internal error") + for _ in range(2): + should_cooldown = _should_cooldown_deployment( + litellm_router_instance=router, + deployment="solo", + exception_status=500, + original_exception=exc, + ) + assert should_cooldown is False, ( + "single-deployment model group must stay protected from a generic allowed_fails override" + ) + + def test_named_exception_policy_still_overrides_single_deployment_protection(self): + """ + Unlike a generic allowed_fails, an explicit per-exception-type allowed_fails_policy + entry is a deliberate, unambiguous opt-in and must still apply even on a + single-deployment model group. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": { + "id": "solo", + "allowed_fails_policy": {"RateLimitErrorAllowedFails": 0}, + }, + }, + ], + ) + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + should_cooldown = _should_cooldown_deployment( + litellm_router_instance=router, + deployment="solo", + exception_status=429, + original_exception=exc, + ) + assert should_cooldown is True, "explicit per-exception-type policy must still cool down a solo deployment" + + +class TestShouldCooldownBasedOnAllowedFailsPolicyFalsyZero: + def test_router_level_policy_of_zero_is_not_swallowed_by_allowed_fails(self): + """ + Router.get_allowed_fails_from_policy returning 0 (a legitimate "cooldown after + the very first failure" policy) must not be treated as falsy and replaced by + router.allowed_fails. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": {"id": "primary"}, + }, + ], + allowed_fails=10, + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + ) + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + should_cooldown = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="primary", + original_exception=exc, + ) + assert should_cooldown is True, "RateLimitErrorAllowedFails=0 must cool down after the first failure" + + +class TestResolveAllowedFailsFromPolicyFallsThrough: + def test_none_value_on_first_match_falls_through_to_next_type(self): + """ + ContentPolicyViolationError is also a BadRequestError; if the policy names + ContentPolicyViolationError but leaves its value unset (None) while setting + BadRequestErrorAllowedFails, resolution must fall through to the + BadRequestError entry rather than stopping at the first isinstance match. + """ + policy = { + "ContentPolicyViolationErrorAllowedFails": None, + "BadRequestErrorAllowedFails": 3, + } + exc = litellm.ContentPolicyViolationError("flagged", "openai", "gpt-4") + result = _resolve_allowed_fails_from_policy(policy=policy, exception=exc) + assert result == 3, "must fall through to BadRequestErrorAllowedFails when the more specific field is unset" + + +class TestDeploymentCallbackOnFailureCooldownTimePrecedence: + def test_model_info_cooldown_time_used_in_primary_sync_path(self): + """ + Router.deployment_callback_on_failure (the primary sync failure-callback path, + as opposed to the fallback path covered by TestFallbackDeploymentCooldown) must + also honor a model_info.cooldown_time, not just litellm_params.cooldown_time. + """ + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4"}, + "model_info": {"id": "primary", "cooldown_time": 15.0}, + }, + ], + ) + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + kwargs = { + "exception": exc, + "litellm_params": { + "model_info": {"id": "primary", "cooldown_time": 15.0}, + }, + } + + with patch("litellm.router._set_cooldown_deployments") as mock_set_cooldown: + router.deployment_callback_on_failure( + kwargs=kwargs, + completion_response=None, + start_time=0, + end_time=1, + ) + + mock_set_cooldown.assert_called_once() + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["time_to_cooldown"] == 15.0, ( + "model_info.cooldown_time must be honored in the primary sync failure-callback path" + ) + + def test_litellm_params_cooldown_time_still_honored_as_fallback(self): + """cooldown_time has pre-existing litellm_params support on this primary + path; it must keep working when model_info doesn't set it.""" + router = _make_router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "openai/gpt-4", "cooldown_time": 20.0}, + "model_info": {"id": "primary"}, + }, + ], + ) + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + kwargs = { + "exception": exc, + "litellm_params": { + "model_info": {"id": "primary"}, + "cooldown_time": 20.0, + }, + } + + with patch("litellm.router._set_cooldown_deployments") as mock_set_cooldown: + router.deployment_callback_on_failure( + kwargs=kwargs, + completion_response=None, + start_time=0, + end_time=1, + ) + + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["time_to_cooldown"] == 20.0, "litellm_params.cooldown_time must still be honored" + + +class TestNewAllowedFailsPolicyFields: + def test_service_unavailable_error_matched_by_policy(self): + """ + ServiceUnavailableError must be matched against ServiceUnavailableErrorAllowedFails. + """ + policy = {"ServiceUnavailableErrorAllowedFails": 0} + exc = litellm.ServiceUnavailableError("Service unavailable", "openai", "gpt-4") + result = _resolve_allowed_fails_from_policy(policy=policy, exception=exc) + assert result == 0 + + def test_bad_gateway_error_matched_by_policy(self): + """ + BadGatewayError must be matched against BadGatewayErrorAllowedFails. + """ + policy = {"BadGatewayErrorAllowedFails": 2} + exc = litellm.BadGatewayError("Bad gateway", "openai", "gpt-4") + result = _resolve_allowed_fails_from_policy(policy=policy, exception=exc) + assert result == 2 + + def test_not_found_error_matched_by_policy(self): + """ + NotFoundError must be matched against NotFoundErrorAllowedFails. + """ + policy = {"NotFoundErrorAllowedFails": 1} + exc = litellm.NotFoundError("Not found", "openai", "gpt-4") + result = _resolve_allowed_fails_from_policy(policy=policy, exception=exc) + assert result == 1 + + def test_unknown_exception_type_returns_none(self): + """ + An exception type not in the policy mapping must return None. + """ + policy = {"RateLimitErrorAllowedFails": 0} + exc = ValueError("unexpected error") + result = _resolve_allowed_fails_from_policy(policy=policy, exception=exc) + assert result is None + + def test_allowed_fails_policy_model_accepts_new_fields(self): + """ + AllowedFailsPolicy Pydantic model must accept the three new fields. + """ + policy = AllowedFailsPolicy( + ServiceUnavailableErrorAllowedFails=3, + BadGatewayErrorAllowedFails=2, + NotFoundErrorAllowedFails=1, + ) + assert policy.ServiceUnavailableErrorAllowedFails == 3 + assert policy.BadGatewayErrorAllowedFails == 2 + assert policy.NotFoundErrorAllowedFails == 1 + + +class TestRouterLevelGetAllowedFailsFromPolicy: + """Router.get_allowed_fails_from_policy must handle all AllowedFailsPolicy fields.""" + + def _make_router(self, **policy_kwargs): + return Router( + model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "fake"}}], + allowed_fails_policy=AllowedFailsPolicy(**policy_kwargs), + ) + + def test_internal_server_error_returned(self): + router = self._make_router(InternalServerErrorAllowedFails=7) + exc = litellm.InternalServerError("500 error", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) == 7 + + def test_service_unavailable_error_returned(self): + router = self._make_router(ServiceUnavailableErrorAllowedFails=4) + exc = litellm.ServiceUnavailableError("503 error", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) == 4 + + def test_bad_gateway_error_returned(self): + router = self._make_router(BadGatewayErrorAllowedFails=2) + exc = litellm.BadGatewayError("502 error", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) == 2 + + def test_not_found_error_returned(self): + router = self._make_router(NotFoundErrorAllowedFails=1) + exc = litellm.NotFoundError("404 error", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) == 1 + + def test_unmatched_exception_returns_none(self): + router = self._make_router(InternalServerErrorAllowedFails=5) + exc = litellm.RateLimitError("429", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) is None diff --git a/tests/router_unit_tests/test_router_cooldown_utils.py b/tests/router_unit_tests/test_router_cooldown_utils.py index ea0cd74d877..6bcb0d9bf84 100644 --- a/tests/router_unit_tests/test_router_cooldown_utils.py +++ b/tests/router_unit_tests/test_router_cooldown_utils.py @@ -19,7 +19,9 @@ from litellm.router_utils.cooldown_handlers import ( _should_cooldown_deployment, cast_exception_status_to_int, _is_cooldown_required, + _has_explicit_allowed_fails_policy_for_exception, ) +from litellm.types.router import AllowedFailsPolicy from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, increment_deployment_successes_for_current_minute, @@ -107,6 +109,137 @@ def test_should_run_cooldown_logic(testing_litellm_router): ) +@pytest.fixture +def single_deployment_router(): + """A router with one deployment whose model_info.id is the lookup-able + "dep-1" (unlike `testing_litellm_router`'s top-level "model_id" key, which + is not absorbed into model_info.id and so never resolves via + get_model_info/get_model_group).""" + return Router( + model_list=[ + { + "model_name": "gpt-5-mini", + "litellm_params": {"model": "gpt-5-mini"}, + "model_info": {"id": "dep-1"}, + }, + ] + ) + + +def test_should_run_cooldown_logic_generic_bad_request_excluded_by_default( + single_deployment_router, +): + """A generic BadRequestError/ContentPolicyViolationError (400) is excluded from + cooldown evaluation by _is_cooldown_required when no allowed_fails_policy is + configured for that exception type. This is the pre-existing, intentional + default: a client error is usually not the deployment's fault.""" + exc = litellm.BadRequestError("bad request", "openai", "gpt-5-mini") + assert ( + _should_run_cooldown_logic(single_deployment_router, "dep-1", 400, exc) is False + ) + + +def test_should_run_cooldown_logic_router_level_policy_does_not_override_bad_request_exclusion( + single_deployment_router, +): + """A router-level allowed_fails_policy is a pre-existing, router-wide setting that + predates the per-deployment override feature, so it must keep its existing behavior + and stay subject to the generic 4XX exclusion. Only an explicit deployment-level + policy (an unambiguous per-exception opt-in for that one deployment) overrides it; + see test_should_run_cooldown_logic_explicit_deployment_level_policy_overrides_content_policy_exclusion.""" + exc = litellm.BadRequestError("bad request", "openai", "gpt-5-mini") + single_deployment_router.allowed_fails_policy = AllowedFailsPolicy( + BadRequestErrorAllowedFails=5 + ) + assert ( + _should_run_cooldown_logic(single_deployment_router, "dep-1", 400, exc) is False + ) + + +def test_should_run_cooldown_logic_explicit_deployment_level_policy_overrides_content_policy_exclusion( + single_deployment_router, +): + """Same as the router-level case, but for a deployment-level allowed_fails_policy + entry (this PR's per-deployment feature) targeting ContentPolicyViolationError.""" + exc = litellm.ContentPolicyViolationError("flagged content", "openai", "gpt-5-mini") + deployment_dict = single_deployment_router.get_model_info(id="dep-1") + deployment_dict["model_info"]["allowed_fails_policy"] = { + "ContentPolicyViolationErrorAllowedFails": 0 + } + assert ( + _should_run_cooldown_logic(single_deployment_router, "dep-1", 400, exc) is True + ) + + +class TestHasExplicitAllowedFailsPolicyForException: + def test_no_policy_anywhere_returns_false(self, single_deployment_router): + exc = litellm.BadRequestError("bad request", "openai", "gpt-5-mini") + assert ( + _has_explicit_allowed_fails_policy_for_exception( + single_deployment_router, "dep-1", exc + ) + is False + ) + + def test_router_level_policy_for_matching_exception_returns_false( + self, single_deployment_router + ): + """Deliberately scoped to deployment-level only: a router-level policy + predates this feature and must not be treated as an explicit per-exception + opt-in for cooldown-gate purposes.""" + exc = litellm.RateLimitError("rate limited", "openai", "gpt-5-mini") + single_deployment_router.allowed_fails_policy = AllowedFailsPolicy( + RateLimitErrorAllowedFails=3 + ) + assert ( + _has_explicit_allowed_fails_policy_for_exception( + single_deployment_router, "dep-1", exc + ) + is False + ) + + def test_router_level_policy_for_different_exception_returns_false( + self, single_deployment_router + ): + exc = litellm.BadRequestError("bad request", "openai", "gpt-5-mini") + single_deployment_router.allowed_fails_policy = AllowedFailsPolicy( + RateLimitErrorAllowedFails=3 + ) + assert ( + _has_explicit_allowed_fails_policy_for_exception( + single_deployment_router, "dep-1", exc + ) + is False + ) + + def test_deployment_level_policy_for_matching_exception_returns_true( + self, single_deployment_router + ): + exc = litellm.ContentPolicyViolationError("flagged", "openai", "gpt-5-mini") + deployment_dict = single_deployment_router.get_model_info(id="dep-1") + deployment_dict["model_info"]["allowed_fails_policy"] = { + "ContentPolicyViolationErrorAllowedFails": 0 + } + assert ( + _has_explicit_allowed_fails_policy_for_exception( + single_deployment_router, "dep-1", exc + ) + is True + ) + + def test_none_deployment_returns_false(self, single_deployment_router): + exc = litellm.RateLimitError("rate limited", "openai", "gpt-5-mini") + single_deployment_router.allowed_fails_policy = AllowedFailsPolicy( + RateLimitErrorAllowedFails=3 + ) + assert ( + _has_explicit_allowed_fails_policy_for_exception( + single_deployment_router, None, exc + ) + is False + ) + + def test_should_cooldown_deployment_rate_limit_error(testing_litellm_router): """ Test the _should_cooldown_deployment function when a rate limit error occurs diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index 1b60c97b510..ebd33aa2e53 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -356,3 +356,142 @@ async def test_ensure_batch_response_returns_early_without_auth(): assert response.output_file_id == "file-raw-output" mock_managed_files.get_unified_output_file_id.assert_not_called() + + +def _in_memory_managed_files(): + """Build a real _PROXY_LiteLLMManagedFiles whose prisma upsert hits an in-memory row.""" + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + + store: dict = {} + + async def _upsert(where, data): + key = where["unified_object_id"] + if key in store: + store[key].update(data["update"]) + else: + store[key] = dict(data["create"]) + + table = MagicMock() + table.upsert = AsyncMock(side_effect=_upsert) + prisma = MagicMock() + prisma.db.litellm_managedobjecttable = table + + cache = MagicMock() + cache.async_set_cache = AsyncMock() + + return ( + _PROXY_LiteLLMManagedFiles(internal_usage_cache=cache, prisma_client=prisma), + store, + ) + + +@pytest.mark.asyncio +async def test_store_unified_object_id_persists_key_and_tags_on_create(): + """Regression (spend loss): the batch create persists the creating key hash and tags so + CheckBatchCost can write an attributed spend row instead of a blank one the DB drops.""" + instance, store = _in_memory_managed_files() + creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + request_tags=["env:prod"], + persist_attribution=True, + ) + + row = store["unified-b"] + assert row["api_key"] == "hash-alice" + assert row["created_by"] == "alice" + assert row["team_id"] == "team-alpha" + assert row["request_tags"].data == ["env:prod"] + + +@pytest.mark.asyncio +async def test_store_unified_object_id_omits_key_and_tags_without_persist_attribution(): + """Regression (spend redirect): a caller that is not the batch create (a poll, or the + generic post-call hook on a retrieve) carries a real hashed key, but must never have it + recorded as the batch's paying key. created_by/team_id keep their existing behavior.""" + instance, store = _in_memory_managed_files() + poller = UserAPIKeyAuth(user_id="bob", team_id="team-bravo", api_key="hash-bob") + + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="in_progress"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=poller, + request_tags=["env:dev"], + ) + + row = store["unified-b"] + assert "api_key" not in row + assert "request_tags" not in row + assert row["created_by"] == "bob" + + +@pytest.mark.asyncio +async def test_store_unified_object_id_attribution_columns_are_write_once(): + """Identity is written only in the upsert create branch, so a later store for the same + batch (a status update, a poll) can neither reassign the paying key nor clear it.""" + instance, store = _in_memory_managed_files() + creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + poller = UserAPIKeyAuth(user_id="bob", team_id="team-bravo", api_key="hash-bob") + + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + request_tags=["env:prod"], + persist_attribution=True, + ) + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="completed"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=poller, + request_tags=["poller-tag"], + persist_attribution=True, + ) + + row = store["unified-b"] + assert row["api_key"] == "hash-alice" + assert row["created_by"] == "alice" + assert row["status"] == "completed" + + upsert_data = instance.prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"] + assert "api_key" not in upsert_data["update"] + assert "request_tags" not in upsert_data["update"] + + +@pytest.mark.asyncio +async def test_store_unified_object_id_omits_unset_columns(): + """A batch created with no tags (the common case) still registers: the optional columns + are omitted rather than passed as None, which prisma rejects for the Json column.""" + instance, store = _in_memory_managed_files() + creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key=None) + + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + request_tags=None, + persist_attribution=True, + ) + + create_data = instance.prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"]["create"] + assert "api_key" not in create_data + assert "request_tags" not in create_data + assert "unified-b" in store diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 63646ae53f8..34cd0cabc2c 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -510,6 +510,117 @@ def _make_real_managed_files_instance(): ) +def _make_object_store_instance(): + """A real store_unified_object_id over an AsyncMock prisma client, so both the + upsert and the update-only write path can be asserted.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + mock_cache = MagicMock() + mock_cache.async_set_cache = AsyncMock() + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + + return ( + _PROXY_LiteLLMManagedFiles( + internal_usage_cache=mock_cache, + prisma_client=mock_prisma, + ), + mock_prisma, + ) + + +@pytest.mark.asyncio +async def test_poll_refreshes_batch_state_without_claiming_the_row(): + """Regression (stale batch state): a poll observes a batch it did not create, so it + must still refresh status and file_object -- otherwise GET /v1/batches serves the + create-time snapshot forever -- while writing none of the attribution columns and + never creating a row it would then own.""" + managed_files, mock_prisma = _make_object_store_instance() + poller = UserAPIKeyAuth( + api_key="sk-the-poller", user_id="bob", team_id="team-bravo", parent_otel_span=None + ) + + await managed_files.store_unified_object_id( + unified_object_id="uoi-1", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="batch-123", + file_purpose="batch", + user_api_key_dict=poller, + request_tags=("poller:tag",), + persist_attribution=False, + create_if_missing=False, + ) + + # the row is refreshed in place, and cannot be conjured by a poll + mock_prisma.db.litellm_managedobjecttable.upsert.assert_not_awaited() + update_many = mock_prisma.db.litellm_managedobjecttable.update_many + update_many.assert_awaited_once() + call = update_many.await_args + assert call.kwargs["where"] == {"unified_object_id": "uoi-1"} + + written = call.kwargs["data"] + assert written["status"] == "completed" + assert json.loads(written["file_object"])["output_file_id"] == "file-output-abc" + # nothing the poller could be billed for + for owned in ("api_key", "request_tags", "created_by", "team_id"): + assert owned not in written + + +@pytest.mark.asyncio +async def test_create_still_upserts_and_claims_attribution(): + """The create is the one caller that can speak for the batch, so it keeps the upsert + (creating the row when absent) and writes the attribution columns.""" + managed_files, mock_prisma = _make_object_store_instance() + creator = UserAPIKeyAuth( + api_key="sk-the-creator", user_id="alice", team_id="team-alpha", parent_otel_span=None + ) + + await managed_files.store_unified_object_id( + unified_object_id="uoi-2", + file_object=_make_batch_response(status="validating"), + litellm_parent_otel_span=None, + model_object_id="batch-456", + file_purpose="batch", + user_api_key_dict=creator, + request_tags=("env:prod",), + persist_attribution=True, + ) + + mock_prisma.db.litellm_managedobjecttable.update_many.assert_not_awaited() + upsert = mock_prisma.db.litellm_managedobjecttable.upsert + upsert.assert_awaited_once() + created = upsert.await_args.kwargs["data"]["create"] + # UserAPIKeyAuth hashes an sk- token on construction; the hash is what is billed + assert created["api_key"] == creator.api_key + assert created["api_key"] != "sk-the-creator" + assert created["created_by"] == "alice" + assert created["team_id"] == "team-alpha" + + +@pytest.mark.asyncio +async def test_default_callers_still_create_their_rows(): + """create_if_missing defaults to True, so the fine-tune, Responses and Anthropic + callers, none of which pass it, keep upserting exactly as before.""" + managed_files, mock_prisma = _make_object_store_instance() + + await managed_files.store_unified_object_id( + unified_object_id="uoi-3", + file_object=_make_batch_response(), + litellm_parent_otel_span=None, + model_object_id="ft-789", + file_purpose="fine-tune", + user_api_key_dict=_make_user_api_key_dict(), + ) + + mock_prisma.db.litellm_managedobjecttable.upsert.assert_awaited_once() + mock_prisma.db.litellm_managedobjecttable.update_many.assert_not_awaited() + + @pytest.mark.asyncio async def test_store_unified_file_id_is_idempotent_via_upsert(): """Regression test for the managed-batch retrieve 500 (UniqueViolationError on diff --git a/tests/test_litellm/integrations/arize/test_arize_utils.py b/tests/test_litellm/integrations/arize/test_arize_utils.py index 83c3351319a..b02fe35cad0 100644 --- a/tests/test_litellm/integrations/arize/test_arize_utils.py +++ b/tests/test_litellm/integrations/arize/test_arize_utils.py @@ -1193,3 +1193,326 @@ def test_arize_coerce_response_obj_returns_original_on_bad_json(): obj = BadJson() assert _coerce_response_obj_for_attrs(obj) is obj + + +def test_arize_mcp_call_tool_result_does_not_break_attribute_setting(): + """`call_mcp_tool` logs the MCP SDK's `CallToolResult`, a Pydantic model + with no `.get`. It used to raise inside `_set_request_attributes`, aborting + the whole attribute block (input messages, invocation params, outputs).""" + from unittest.mock import MagicMock + + from mcp.types import CallToolResult, TextContent + + span = MagicMock() + kwargs = { + "model": "MCP: get_weather", + "standard_logging_object": { + "model_parameters": {"user": "u-1"}, + "metadata": {}, + "call_type": "call_mcp_tool", + }, + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "mcp"}, + } + response_obj = CallToolResult( + content=[TextContent(type="text", text="sunny, 21C")], isError=False + ) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + span.record_exception.assert_not_called() + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.OPENINFERENCE_SPAN_KIND] == "TOOL" + assert written["llm.request.type"] == "call_mcp_tool" + # Emitted after the old crash point, so absent before the fix. + assert written[SpanAttributes.LLM_INVOCATION_PARAMETERS] == '{"user": "u-1"}' + assert written[SpanAttributes.USER_ID] == "u-1" + + +def test_arize_coerce_response_obj_dumps_pydantic_without_get(): + from mcp.types import CallToolResult, TextContent + + from litellm.integrations.arize._utils import _coerce_response_obj_for_attrs + + result = CallToolResult(content=[TextContent(type="text", text="hi")], isError=False) + coerced = _coerce_response_obj_for_attrs(result) + + assert isinstance(coerced, dict) + assert coerced["isError"] is False + assert coerced["content"][0]["text"] == "hi" + + +def test_arize_request_attributes_survive_uncoercible_response_obj(): + """A response object that is neither dict-like nor coercible (binary + passthrough body, SDK object) must not abort attribute setting.""" + from unittest.mock import MagicMock + + span = MagicMock() + kwargs = { + "model": "gpt-4o", + "standard_logging_object": { + "model_parameters": {}, + "metadata": {}, + "call_type": "completion", + }, + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + } + + class Opaque: + pass + + ArizeLogger.set_arize_attributes(span, kwargs, Opaque()) + + span.record_exception.assert_not_called() + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written["llm.provider"] == "openai" + + +def _mcp_kwargs(mcp_tool_call_metadata=None, **overrides): + return { + "model": "MCP: get_weather", + "standard_logging_object": { + "model_parameters": {}, + "metadata": { + "mcp_tool_call_metadata": mcp_tool_call_metadata + or { + "name": "get_weather", + "arguments": {"city": "Seoul"}, + "namespaced_tool_name": "weather-mcp/get_weather", + } + }, + "call_type": "call_mcp_tool", + }, + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "mcp"}, + **overrides, + } + + +def test_arize_mcp_tool_span_renders_name_input_and_output(): + """`call_mcp_tool` spans have no messages/choices, so Input and Output came + out blank. Render them from mcp_tool_call_metadata + CallToolResult.""" + from unittest.mock import MagicMock + + from mcp.types import CallToolResult, TextContent + + span = MagicMock() + response_obj = CallToolResult( + content=[TextContent(type="text", text="sunny, 21C")], isError=False + ) + + ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.TOOL_NAME] == "get_weather" + assert written[SpanAttributes.INPUT_VALUE] == '{"city": "Seoul"}' + assert written[SpanAttributes.INPUT_MIME_TYPE] == "application/json" + assert written[SpanAttributes.OUTPUT_VALUE] == "sunny, 21C" + assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "text/plain" + + +def test_arize_mcp_tool_span_serializes_non_text_content(): + """Image/resource results have no text part, so fall back to JSON.""" + from unittest.mock import MagicMock + + from mcp.types import CallToolResult, ImageContent + + span = MagicMock() + response_obj = CallToolResult( + content=[ImageContent(type="image", data="Zm9v", mimeType="image/png")], + isError=False, + ) + + ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json" + assert "image/png" in written[SpanAttributes.OUTPUT_VALUE] + + +def test_arize_mcp_tool_span_respects_message_redaction(): + """Tool arguments and results are user content. With redaction on, only the + tool name may reach the span.""" + from unittest.mock import MagicMock + + from mcp.types import CallToolResult, TextContent + + span = MagicMock() + response_obj = CallToolResult( + content=[TextContent(type="text", text="SSN 123-45-6789")], isError=False + ) + + ArizeLogger.set_arize_attributes( + span, + _mcp_kwargs(standard_callback_dynamic_params={"turn_off_message_logging": True}), + response_obj, + ) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.TOOL_NAME] == "get_weather" + assert SpanAttributes.INPUT_VALUE not in written + assert SpanAttributes.OUTPUT_VALUE not in written + + +def test_arize_non_mcp_span_gets_no_tool_name(): + """The MCP emitter must not fire on ordinary completions.""" + from unittest.mock import MagicMock + + from litellm.types.utils import Choices, ModelResponse + + span = MagicMock() + kwargs = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "standard_logging_object": { + "model_parameters": {}, + "metadata": {"mcp_tool_call_metadata": {"name": "get_weather"}}, + "call_type": "completion", + }, + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "openai"}, + } + response_obj = ModelResponse( + choices=[Choices(message={"role": "assistant", "content": "hello"})], + model="gpt-4o", + id="r-1", + ) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert SpanAttributes.TOOL_NAME not in written + assert written[SpanAttributes.OUTPUT_VALUE] == "hello" + + +def test_arize_mcp_tool_span_renders_empty_arguments(): + from unittest.mock import MagicMock + + from mcp.types import CallToolResult, TextContent + + span = MagicMock() + kwargs = _mcp_kwargs(mcp_tool_call_metadata={"name": "ping", "arguments": {}}) + response_obj = CallToolResult(content=[TextContent(type="text", text="pong")], isError=False) + + ArizeLogger.set_arize_attributes(span, kwargs, response_obj) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.INPUT_VALUE] == "{}" + assert written[SpanAttributes.INPUT_MIME_TYPE] == "application/json" + + +def test_arize_mcp_tool_span_renders_empty_content(): + from unittest.mock import MagicMock + + from mcp.types import CallToolResult + + span = MagicMock() + response_obj = CallToolResult(content=[], isError=False) + + ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.OUTPUT_VALUE] == "[]" + assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json" + + +def test_arize_mcp_tool_span_falls_back_to_structured_content(): + from unittest.mock import MagicMock + + from mcp.types import CallToolResult + + span = MagicMock() + response_obj = CallToolResult(content=[], structuredContent={"temp_c": 21}, isError=False) + + ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.OUTPUT_VALUE] == '{"temp_c": 21}' + assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json" + + +def test_arize_list_mcp_tools_response_does_not_break_attribute_setting(): + from unittest.mock import MagicMock + + span = MagicMock() + kwargs = { + "model": "MCP: list_tools", + "messages": [{"role": "user", "content": "list"}], + "standard_logging_object": { + "model_parameters": {}, + "metadata": {}, + "call_type": "list_mcp_tools", + }, + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "mcp"}, + } + + ArizeLogger.set_arize_attributes(span, kwargs, [{"name": "get_weather"}]) + + span.record_exception.assert_not_called() + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written["llm.input_messages.0.message.content"] == "list" + + +def test_arize_mcp_tool_span_serializes_mixed_text_and_media(): + from unittest.mock import MagicMock + + from mcp.types import CallToolResult, ImageContent, TextContent + + span = MagicMock() + response_obj = CallToolResult( + content=[ + TextContent(type="text", text="see image"), + ImageContent(type="image", data="Zm9v", mimeType="image/png"), + ], + isError=False, + ) + + ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), response_obj) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json" + assert "see image" in written[SpanAttributes.OUTPUT_VALUE] + assert "image/png" in written[SpanAttributes.OUTPUT_VALUE] + + +def test_arize_mcp_tool_span_without_response_object_keeps_name_and_input(): + from unittest.mock import MagicMock + + span = MagicMock() + + ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), None) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.TOOL_NAME] == "get_weather" + assert written[SpanAttributes.INPUT_VALUE] == '{"city": "Seoul"}' + assert SpanAttributes.OUTPUT_VALUE not in written + + +def test_arize_mcp_tool_span_without_content_emits_no_output(): + from unittest.mock import MagicMock + + span = MagicMock() + + ArizeLogger.set_arize_attributes(span, _mcp_kwargs(), {"isError": False}) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert written[SpanAttributes.TOOL_NAME] == "get_weather" + assert SpanAttributes.OUTPUT_VALUE not in written + + +def test_arize_mcp_emitter_is_inert_without_a_standard_logging_object(): + from unittest.mock import MagicMock + + span = MagicMock() + kwargs = { + "model": "MCP: get_weather", + "optional_params": {}, + "litellm_params": {"custom_llm_provider": "mcp"}, + } + + ArizeLogger.set_arize_attributes(span, kwargs, None) + + written = {c.args[0]: c.args[1] for c in span.set_attribute.call_args_list} + assert SpanAttributes.TOOL_NAME not in written 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 0970e029ad8..158cdb45f6b 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 @@ -2143,6 +2143,54 @@ def test_token_type_cost_breakdown_matches_real_gemini_numbers(): assert breakdown.cache_creation_cost == 0.0 +def test_token_type_cost_breakdown_xai_at_exactly_200k_uses_higher_tier_rates(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + usage = Usage( + prompt_tokens=200_000, + completion_tokens=2_000, + total_tokens=202_000, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=1_500, text_tokens=500 + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=50_000, text_tokens=150_000 + ), + ) + + breakdown = get_token_type_cost_breakdown( + model="grok-4.20-0309-reasoning", custom_llm_provider="xai", usage=usage + ) + + assert breakdown.reasoning_cost == pytest.approx(1_500 * 5e-06) + assert breakdown.cache_read_cost == pytest.approx(50_000 * 4e-07) + + +def test_token_type_cost_breakdown_xai_just_below_200k_uses_base_tier_rates(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + usage = Usage( + prompt_tokens=199_999, + completion_tokens=2_000, + total_tokens=201_999, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=1_500, text_tokens=500 + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=50_000, text_tokens=149_999 + ), + ) + + breakdown = get_token_type_cost_breakdown( + model="grok-4.20-0309-reasoning", custom_llm_provider="xai", usage=usage + ) + + assert breakdown.reasoning_cost == pytest.approx(1_500 * 2.5e-06) + assert breakdown.cache_read_cost == pytest.approx(50_000 * 2e-07) + + def test_token_type_cost_breakdown_includes_cache_creation_from_top_level_usage(): """ Bedrock/Anthropic report cache tokens as top-level usage fields; the Usage @@ -2446,6 +2494,60 @@ def test_token_type_cost_breakdown_applies_regional_uplift(): assert text_input_cost + eu.cache_read_cost == pytest.approx(prompt_cost) +@pytest.mark.parametrize("details_as_dict", [True, False]) +def test_image_response_input_image_tokens_priced_at_image_rate(details_as_dict): + """ + Image input tokens must be priced at input_cost_per_image_token even when + input_tokens_details is a plain dict, as in OpenAI image edit responses. + + Regression test: dict-shaped input_tokens_details was read with getattr(), + which returns None for dicts, so image input tokens silently fell back to + the text input rate (e.g. $5/M instead of $8/M for gpt-image-2). + """ + from unittest.mock import patch + + from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_image_response_cost_from_usage, + ) + from litellm.types.utils import Usage + + mock_model_info = { + "input_cost_per_token": 5e-6, + "input_cost_per_image_token": 8e-6, + "output_cost_per_image_token": 3e-5, + } + + input_details = {"text_tokens": 19, "image_tokens": 512} + image_response = ImageResponse(data=[ImageObject(b64_json="x")]) + # Mirror the usage shape of a real OpenAI images.edit response: + # a Usage object carrying input_tokens/output_tokens with detail dicts. + image_response.usage = Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=689, + input_tokens=531, + input_tokens_details=( + input_details + if details_as_dict + else ImageUsageInputTokensDetails(**input_details) + ), + output_tokens=158, + output_tokens_details={"image_tokens": 158, "text_tokens": 0}, + ) + + with patch( + "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", + return_value=mock_model_info, + ): + cost = calculate_image_response_cost_from_usage( + model="gpt-image-2", + image_response=image_response, + custom_llm_provider="openai", + ) + + expected = 19 * 5e-6 + 512 * 8e-6 + 158 * 3e-5 + assert cost is not None + assert round(cost, 12) == round(expected, 12) GEMINI_DAY0_LAUNCH_PRICING = [ ("gemini-3.6-flash", 1.5e-06, 7.5e-06, 1.5e-07), ("gemini/gemini-3.6-flash", 1.5e-06, 7.5e-06, 1.5e-07), 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 df55e1ee433..d83a721aa55 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -15,6 +15,7 @@ import httpx from openai._legacy_response import HttpxBinaryResponseContent import litellm +from litellm._logging import session_id_var, trace_id_var from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging @@ -3392,6 +3393,51 @@ def test_failure_handler_runs_sync_callbacks_for_non_pass_through_requests( dummy_logger.log_failure_event.assert_called_once() +@pytest.mark.asyncio +async def test_async_failure_handler_runs_callbacks_and_restores_correlation_context(logging_obj): + """await logging_obj.async_failure_handler(...) must dispatch async failure callbacks + and, once its own body completes, restore trace_id/session_id contextvars via + _restore_correlation_context() (the fix for the nested-call context leak).""" + from litellm._logging import session_id_var, trace_id_var + from litellm.integrations.custom_logger import CustomLogger + + class DummyLogger(CustomLogger): + pass + + logging_obj.call_type = "acompletion" + logging_obj.stream = False + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + + dummy_logger = DummyLogger() + dummy_logger.async_log_failure_event = AsyncMock() + + # logging_obj is constructed by the fixture (before this line runs), so it + # already captured whatever was ambient at that point as its own pre-call + # value - assert restoration lands back on THAT captured value, not a + # value set here (which would be too late to affect __init__'s snapshot). + trace_id_var.set("mutated-during-call") + session_id_var.set("mutated-during-call") + try: + with patch.object( + logging_obj, + "get_combined_callback_list", + return_value=[dummy_logger], + ): + await logging_obj.async_failure_handler( + exception=Exception("test error"), + traceback_exception="", + ) + + dummy_logger.async_log_failure_event.assert_called_once() + assert trace_id_var.get() == logging_obj._pre_call_trace_id + assert session_id_var.get() == logging_obj._pre_call_session_id + assert trace_id_var.get() != "mutated-during-call" + finally: + trace_id_var.set("") + session_id_var.set("") + + def test_merge_hidden_params_from_response_into_metadata_populates_metadata(): """Streaming completion path should mirror non-stream: metadata.hidden_params from response.""" from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -4310,3 +4356,199 @@ def test_pre_call_does_not_pin_request_in_module_state(logging_obj): logging_obj.post_call(original_response='{"ok": true}', input=big_input, api_key="sk-test") assert litellm.error_logs == {} + + +def test_logging_init_sets_trace_id(): + """Logging.__init__() must call set_trace_id with self.litellm_trace_id.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + trace_id_var.set("") + + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-001", + function_id="fn-001", + kwargs={}, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + + +def test_logging_init_skips_stamping_when_correlation_logging_unsupported(): + """supports_correlation_logging=False (what wrapper(), the sync entry + point, always passes) must leave trace_id_var/session_id_var completely + untouched, even though self.litellm_trace_id/litellm_session_id (the + plain attributes used by StandardLoggingPayload) are still populated as + usual - only the ambient contextvar stamping is gated.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + trace_id_var.set("") + session_id_var.set("") + + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-sync-excluded", + function_id="fn-sync-excluded", + kwargs={"litellm_session_id": "should-not-be-stamped"}, + litellm_trace_id="should-not-be-stamped-either", + supports_correlation_logging=False, + ) + + assert trace_id_var.get() == "" + assert session_id_var.get() == "" + # The plain attributes are unaffected - only the contextvar stamping is gated. + assert log_obj.litellm_trace_id == "should-not-be-stamped-either" + assert log_obj.litellm_session_id == "should-not-be-stamped" + + +def test_logging_init_sets_session_id_when_provided(): + """Logging.__init__() must call set_session_id when litellm_session_id is in kwargs.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + session_id_var.set("") + + Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-002", + function_id="fn-002", + kwargs={"litellm_session_id": "my-session-99"}, + ) + assert session_id_var.get() == "my-session-99" + + +def test_logging_init_resets_session_id_to_empty_when_absent(): + """When no session_id is in kwargs, Logging.__init__() must reset session_id_var to "" + so a prior request's session_id does not leak into subsequent log records.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + session_id_var.set("preexisting-sid") + + Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-003", + function_id="fn-003", + kwargs={}, + ) + assert session_id_var.get() == "" + + +def test_restore_correlation_context_resets_to_pre_call_value(): + """_restore_correlation_context() must put trace_id_var/session_id_var back to + whatever they were immediately before this Logging instance was constructed. + This is the mechanism that prevents a nested call (e.g. a guardrail's own + LLM-as-judge call sharing the same asyncio Task) from leaking its trace_id/ + session_id into the outer call's subsequent log lines.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + trace_id_var.set("outer-trace") + session_id_var.set("outer-session") + try: + inner = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="inner-call", + function_id="fn-inner", + kwargs={"litellm_session_id": "inner-session"}, + ) + assert trace_id_var.get() == inner.litellm_trace_id + assert session_id_var.get() == "inner-session" + + inner._restore_correlation_context() + + assert trace_id_var.get() == "outer-trace" + assert session_id_var.get() == "outer-session" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_restore_correlation_context_safe_to_call_repeatedly(): + """Calling _restore_correlation_context() more than once must not raise. + + It's deliberately NOT guarded against repeat calls: wrapper()'s finally + block and a terminal handler (success_handler/failure_handler) can both + end up calling it for the same instance, potentially from different + asyncio Tasks - each call needs to take effect in its own Task's view of + the contextvars, so repeat calls are expected, not just tolerated.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="call-idempotent", + function_id="fn-idempotent", + kwargs={}, + ) + log_obj._restore_correlation_context() + log_obj._restore_correlation_context() # must not raise + + +@pytest.mark.asyncio +async def test_restore_correlation_context_works_across_asyncio_task_boundary(): + """_restore_correlation_context() must succeed even when it's called from a + different asyncio Task than the one Logging.__init__() ran in - exactly what + happens on litellm's real async success path, where async_success_handler is + dispatched via asyncio.create_task / the global logging worker rather than + awaited directly in the request's own task. + + A contextvars.Token can only be reset in the exact Context it was created in + and raises ValueError otherwise (verified separately against raw contextvars, + not just this codebase). The fix uses a plain set() of the captured pre-call + value instead, which works regardless of which Task calls it. This test + fails with a token-based implementation - the child task's reset() would + raise, get silently swallowed, and leave the child's view unrestored - and + passes with the value-based one. + """ + from litellm.litellm_core_utils.litellm_logging import Logging + + trace_id_var.set("outer-trace-cross-task") + session_id_var.set("outer-session-cross-task") + try: + # __init__ runs in THIS (outer) task's context. + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=None, + litellm_call_id="cross-task-call", + function_id="fn-cross-task", + kwargs={"litellm_session_id": "cross-task-session"}, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "cross-task-session" + + async def restore_in_new_task(): + # Simulates async_success_handler running in a task spawned after + # __init__ already ran elsewhere - a different Context object. + log_obj._restore_correlation_context() + return trace_id_var.get(), session_id_var.get() + + trace_in_child, session_in_child = await asyncio.create_task(restore_in_new_task()) + + assert trace_in_child == "outer-trace-cross-task" + assert session_in_child == "outer-session-cross-task" + finally: + trace_id_var.set("") + session_id_var.set("") diff --git a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py index a417ad90eb7..1eb49f4859f 100644 --- a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py +++ b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py @@ -9,7 +9,7 @@ Covers: import os import sys import time -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -69,8 +69,12 @@ class TestCustomStreamWrapperMaxDuration: @pytest.mark.asyncio async def test_should_raise_on_async_anext_when_exceeded(self): - """__anext__ should check the limit before iterating.""" + """__anext__ should check the limit before iterating, dispatching the + same failure-callback/logging path every other stream failure goes + through (dispatch_failure_handlers is async on the real Logging class, + so the mock needs to be awaitable too).""" wrapper = _make_custom_stream_wrapper() + wrapper.logging_obj.dispatch_failure_handlers = AsyncMock() wrapper._stream_created_time = time.time() - 20 with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): with pytest.raises(litellm.Timeout): diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 5806b37539c..101935cac0a 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -14,6 +14,8 @@ import traceback from typing import Optional import litellm +from litellm import verbose_logger +from litellm._logging import session_id_var, trace_id_var from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import ( AUDIO_ATTRIBUTE, @@ -3551,3 +3553,613 @@ def test_openai_custom_tool_call_stream_deltas_survive_conversion(logging_obj: L assert combined_input == "*** Begin Patch\n*** End Patch\n" finish_reasons = [chunk.choices[0].finish_reason for chunk in emitted if chunk.choices] assert "tool_calls" in finish_reasons + + +def test_sync_completion_never_stamps_correlation_context(monkeypatch): + """wrapper() (the sync entry point) does not participate in + request_correlation_in_logs at all: Logging.__init__() is called with + supports_correlation_logging=False for every sync call, so + trace_id_var/session_id_var are never touched, regardless of whether the + caller passes litellm_trace_id/litellm_session_id or the call streams. + + This is a deliberate scoping decision, not an oversight: a plain OS + thread has no per-call isolation the way an asyncio Task does, and a + thread pool's worker threads are recycled across unrelated requests, so + safely supporting this for the sync path needs its own restore mechanism + with its own tests - tracked as a separate, follow-up piece of work. + Async (acompletion/wrapper_async, the only path the proxy uses) is + unaffected - see test_async_streaming_completion_does_not_reset_context_before_iteration.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + # Reset explicitly rather than asserting a clean slate - this must hold + # regardless of what any other test left behind in these module-level + # contextvars. + trace_id_var.set("") + session_id_var.set("") + try: + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + litellm_trace_id="should-never-appear", + litellm_session_id="should-never-appear-either", + num_retries=0, + ) + assert trace_id_var.get() == "" + assert session_id_var.get() == "" + + response = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + stream=True, + litellm_trace_id="should-never-appear-stream", + litellm_session_id="should-never-appear-stream-either", + num_retries=0, + ) + for _ in response: + pass + assert trace_id_var.get() == "" + assert session_id_var.get() == "" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_abandoned_sync_stream_cannot_contaminate_a_later_call_on_the_same_thread(monkeypatch): + """The maintainer-reported blocking bug reproduced live in this session - + request A starts a sync stream, consumes one chunk, abandons it; request + B runs next on the same forced-reuse ThreadPoolExecutor worker - is now + structurally impossible rather than merely restored-after-the-fact: since + sync calls never stamp trace_id_var/session_id_var at all + (supports_correlation_logging=False), there is nothing for request A to + leave behind for request B to inherit.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + from concurrent.futures import ThreadPoolExecutor + + pool = ThreadPoolExecutor(max_workers=1) + try: + + def call_a_abandon_stream(): + response = litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "call A"}], + mock_response="call A response", + stream=True, + litellm_session_id="SESSION-AAA", + litellm_trace_id="TRACE-AAA", + num_retries=0, + ) + next(response) # consume exactly one chunk, then abandon it + + def call_b_non_streaming(): + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "call B"}], + mock_response="call B response", + litellm_session_id="SESSION-BBB", + litellm_trace_id="TRACE-BBB", + num_retries=0, + ) + return trace_id_var.get(), session_id_var.get() + + pool.submit(call_a_abandon_stream).result() + ids_after_b = pool.submit(call_b_non_streaming).result() + + assert ids_after_b == ("", "") + finally: + pool.shutdown(wait=True) + + +@pytest.mark.asyncio +async def test_async_streaming_completion_does_not_reset_context_before_iteration(monkeypatch): + """Same as above for wrapper_async()/acompletion().""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + trace_id_var.set("outer-trace-async-stream") + session_id_var.set("outer-session-async-stream") + try: + response = await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + stream=True, + litellm_session_id="async-streaming-call-session", + num_retries=0, + ) + assert session_id_var.get() == "async-streaming-call-session" + + async for _ in response: + pass + + # Once the stream is genuinely exhausted, the *consuming* task's own + # context must be restored - async_success_handler's own dispatch (via + # asyncio.create_task) only fixes up its own detached task, not this one. + assert session_id_var.get() == "outer-session-async-stream" + assert trace_id_var.get() == "outer-trace-async-stream" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_del_restores_correlation_context(): + """CustomStreamWrapper.__del__ is the best-effort fallback for an abandoned + stream (caller never exhausts it, so the normal terminal-handler restore + never fires). Testing this via real garbage collection is unreliable in + practice - CPython's per-chunk logging submits work to a thread pool + executor whose worker thread transiently holds its own reference to the + wrapper (a bound method argument) until that task completes, so refcount + doesn't reliably hit zero on a deterministic schedule even with polling. + Call __del__ directly instead: it's a plain method, calling it early + doesn't run actual finalization, and this exercises exactly the logic that + real garbage collection would eventually trigger. + """ + trace_id_var.set("outer-trace-abandoned") + session_id_var.set("outer-session-abandoned") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="abandoned-stream-call", + function_id="fn-abandoned-stream", + kwargs={"litellm_session_id": "abandoned-stream-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + wrapper.__del__() + + assert trace_id_var.get() == "outer-trace-abandoned" + assert session_id_var.get() == "outer-session-abandoned" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_del_never_raises_with_broken_logging_obj(): + """__del__ runs during garbage collection, possibly at interpreter + shutdown - it must never raise regardless of what's wrong with logging_obj, + or Python prints an ignored "exception in __del__" warning and, worse, + could mask the real error a caller is in the middle of handling.""" + + class ExplodingLogging: + model_call_details: dict = {} + + def _restore_correlation_context(self): + raise RuntimeError("logging_obj is in a bad state") + + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=ExplodingLogging(), + ) + wrapper.__del__() # must not raise + + +def test_stream_wrapper_del_does_not_clobber_a_newer_active_call(): + """A delayed finalizer must never stomp a different, still-active call's + context. If an abandoned stream's __del__ fires late - after a new call + has already started in the same Task/thread and claimed the contextvars - + unconditionally restoring the abandoned stream's own pre-call snapshot + would corrupt the active call's subsequent log lines with stale ids.""" + trace_id_var.set("outer-trace-before-abandoned-call") + session_id_var.set("outer-session-before-abandoned-call") + try: + abandoned_log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="abandoned-stream-call", + function_id="fn-abandoned-stream", + kwargs={"litellm_session_id": "abandoned-stream-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=abandoned_log_obj, + ) + + # A new, unrelated call starts in this same Task/thread before the + # abandoned stream's __del__ ever fires, and claims the contextvars. + Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="newer-active-call", + function_id="fn-newer-active-call", + kwargs={"litellm_session_id": "newer-active-session"}, + ) + assert trace_id_var.get() != abandoned_log_obj.litellm_trace_id + assert session_id_var.get() == "newer-active-session" + + # The delayed finalizer for the abandoned stream must not clobber + # the newer call's still-active ids. + wrapper.__del__() + + assert trace_id_var.get() != abandoned_log_obj.litellm_trace_id + assert session_id_var.get() == "newer-active-session" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_del_restores_when_own_session_id_needed_sanitizing(): + """The __del__ guard must compare against the *sanitized* id actually + stored in the contextvar, not the raw litellm_session_id/litellm_trace_id + - set_session_id()/set_trace_id() strip control characters before + storing, so a caller-supplied id containing e.g. a newline would never + equal the raw attribute, and the guard would wrongly conclude some other + call has claimed the context and skip cleanup forever.""" + trace_id_var.set("outer-trace-needs-sanitizing") + session_id_var.set("outer-session-needs-sanitizing") + try: + raw_session_id = "abandoned\nsession\rwith-control-chars" + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="abandoned-stream-needs-sanitizing", + function_id="fn-abandoned-stream-needs-sanitizing", + kwargs={"litellm_session_id": raw_session_id}, + ) + # Sanity: the contextvar holds the sanitized value, which differs + # from the raw litellm_session_id this test constructed it with. + assert session_id_var.get() != raw_session_id + assert log_obj.litellm_session_id == raw_session_id + + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + wrapper.__del__() + + assert trace_id_var.get() == "outer-trace-needs-sanitizing" + assert session_id_var.get() == "outer-session-needs-sanitizing" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk(): + """When the underlying stream ends without ever emitting an explicit + finish_reason chunk, __next__ synthesizes one via finish_reason_handler() + and returns it. That chunk is still this call's own data - the caller's + own (application-level) log statements processing it run immediately + after this return, in the same synchronous frame, so context must NOT be + restored yet or those log lines would carry the wrong ids. A caller that + keeps iterating (the common, non-early-break pattern) still gets a + correct, deterministic restore on the very next __next__() call, since + completion_stream is already exhausted and immediately re-raises + StopIteration.""" + trace_id_var.set("outer-trace-finish-reason") + session_id_var.set("outer-session-finish-reason") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="finish-reason-call", + function_id="fn-finish-reason", + kwargs={"litellm_session_id": "finish-reason-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "finish-reason-session" + + chunk = next(wrapper) + + assert chunk.choices[0].finish_reason is not None + # Still this call's own ids - not restored yet. + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "finish-reason-session" + + # A caller that keeps iterating (doesn't break early) still gets a + # deterministic restore right here, on the next real StopIteration. + with pytest.raises(StopIteration): + next(wrapper) + assert trace_id_var.get() == "outer-trace-finish-reason" + assert session_id_var.get() == "outer-session-finish-reason" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_stream_wrapper_del_cleans_up_after_synthesized_finish_reason_chunk(): + """A caller that breaks immediately after seeing finish_reason (the + early-break pattern) never triggers the next()-driven restore above - it + relies on the best-effort __del__ guard instead, same as any other + abandoned stream. The guard must still recognize this call's own + (unrestored) ids as unclaimed and clean them up.""" + trace_id_var.set("outer-trace-finish-reason-del") + session_id_var.set("outer-session-finish-reason-del") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="finish-reason-del-call", + function_id="fn-finish-reason-del", + kwargs={"litellm_session_id": "finish-reason-del-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + + chunk = next(wrapper) + assert chunk.choices[0].finish_reason is not None + + wrapper.__del__() + + assert trace_id_var.get() == "outer-trace-finish-reason-del" + assert session_id_var.get() == "outer-session-finish-reason-del" + finally: + trace_id_var.set("") + session_id_var.set("") + + +@pytest.mark.asyncio +async def test_stream_wrapper_anext_keeps_context_active_through_synthesized_finish_reason_chunk(): + """Async sibling of test_stream_wrapper_next_keeps_context_active_through_synthesized_finish_reason_chunk - + _finalize_completed_stream()'s else branch must not restore before + returning the synthesized chunk either.""" + trace_id_var.set("outer-trace-anext-finish-reason") + session_id_var.set("outer-session-anext-finish-reason") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="anext-finish-reason-call", + function_id="fn-anext-finish-reason", + kwargs={"litellm_session_id": "anext-finish-reason-session"}, + ) + + async def _empty_aiter(): + return + yield # pragma: no cover - makes this an async generator + + wrapper = CustomStreamWrapper( + completion_stream=_empty_aiter(), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "anext-finish-reason-session" + + chunk = await wrapper.__anext__() + + assert chunk.choices[0].finish_reason is not None + # Still this call's own ids - not restored yet. + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "anext-finish-reason-session" + + # A caller that keeps iterating still gets a deterministic restore + # right here, on the next real StopAsyncIteration. + with pytest.raises(StopAsyncIteration): + await wrapper.__anext__() + assert trace_id_var.get() == "outer-trace-anext-finish-reason" + assert session_id_var.get() == "outer-session-anext-finish-reason" + finally: + trace_id_var.set("") + session_id_var.set("") + + +@pytest.mark.asyncio +async def test_stream_wrapper_anext_max_duration_timeout_restores_consumer_correlation_context(monkeypatch): + """_check_max_streaming_duration() raises litellm.Timeout when a client keeps + an async stream open past LITELLM_MAX_STREAMING_DURATION_SECONDS. That raise + must flow through the same except Exception -> _handle_stream_fallback_error + path as every other failure so the consumer's outer correlation context gets + restored - calling the check before entering __anext__()'s try block would + let the Timeout bypass that restoration entirely.""" + monkeypatch.setattr(litellm.constants, "LITELLM_MAX_STREAMING_DURATION_SECONDS", 1) + trace_id_var.set("outer-trace-max-duration") + session_id_var.set("outer-session-max-duration") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="max-duration-call", + function_id="fn-max-duration", + kwargs={"litellm_session_id": "max-duration-session"}, + ) + + async def _empty_aiter(): + return + yield # pragma: no cover - makes this an async generator + + wrapper = CustomStreamWrapper( + completion_stream=_empty_aiter(), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "max-duration-session" + + wrapper._stream_created_time = time.time() - 10 + + with pytest.raises(Exception): + await wrapper.__anext__() + + assert trace_id_var.get() == "outer-trace-max-duration" + assert session_id_var.get() == "outer-session-max-duration" + finally: + trace_id_var.set("") + session_id_var.set("") + + +@pytest.mark.asyncio +async def test_stream_wrapper_aclose_restores_consumer_correlation_context(): + """Explicit early termination (aclose(), e.g. on client disconnect or a + router fallback aborting an in-progress stream) must restore the caller's + correlation context too - not just __del__'s best-effort GC-timed fallback, + since aclose() is normally called deterministically by the consumer/ + framework, unlike __del__.""" + trace_id_var.set("outer-trace-aclose") + session_id_var.set("outer-session-aclose") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="aclose-call", + function_id="fn-aclose", + kwargs={"litellm_session_id": "aclose-session"}, + ) + + async def _empty_aiter(): + return + yield # pragma: no cover - makes this an async generator + + wrapper = CustomStreamWrapper( + completion_stream=_empty_aiter(), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "aclose-session" + + await wrapper.aclose() + + assert trace_id_var.get() == "outer-trace-aclose" + assert session_id_var.get() == "outer-session-aclose" + finally: + trace_id_var.set("") + session_id_var.set("") + + +@pytest.mark.asyncio +async def test_stream_wrapper_aclose_keeps_context_active_through_close_failure_diagnostic(monkeypatch): + """If closing the underlying provider stream raises, aclose()'s except + branch logs a debug diagnostic. That log line must still carry the + closing stream's own trace_id/session_id - the outer context must not be + restored until after the close attempt (and its diagnostic) completes.""" + trace_id_var.set("outer-trace-close-fail") + session_id_var.set("outer-session-close-fail") + try: + log_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="close-fail-call", + function_id="fn-close-fail", + kwargs={"litellm_session_id": "close-fail-session"}, + ) + + class _RaisingAsyncCloseStream: + async def aclose(self): + raise RuntimeError("boom closing stream") + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + wrapper = CustomStreamWrapper( + completion_stream=_RaisingAsyncCloseStream(), + model="gpt-3.5-turbo", + logging_obj=log_obj, + ) + assert trace_id_var.get() == log_obj.litellm_trace_id + assert session_id_var.get() == "close-fail-session" + + captured_ids = {} + real_debug = verbose_logger.debug + + def fake_debug(msg, *args, **kwargs): + if "error closing completion_stream" in msg: + captured_ids["trace_id"] = trace_id_var.get() + captured_ids["session_id"] = session_id_var.get() + return real_debug(msg, *args, **kwargs) + + monkeypatch.setattr(verbose_logger, "debug", fake_debug) + + await wrapper.aclose() + + assert captured_ids["trace_id"] == log_obj.litellm_trace_id + assert captured_ids["session_id"] == "close-fail-session" + assert trace_id_var.get() == "outer-trace-close-fail" + assert session_id_var.get() == "outer-session-close-fail" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_handle_stream_fallback_error_restores_context_only_after_exception_mapping(monkeypatch): + """_map_anthropic_exception/_map_aleph_alpha_exception synchronously log a + debug diagnostic (the raw status code) as part of exception_type()'s + mapping. The consumer's outer context must not be restored until that + mapping call returns, or the diagnostic log line would carry the outer + (or empty) trace_id/session_id instead of the failing stream's own.""" + trace_id_var.set("outer-trace-fallback") + session_id_var.set("outer-session-fallback") + try: + log_obj = Logging( + model="claude-3-opus", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=None, + litellm_call_id="fallback-error-call", + function_id="fn-fallback-error", + kwargs={"litellm_session_id": "fallback-error-session"}, + ) + wrapper = CustomStreamWrapper( + completion_stream=iter([]), + model="claude-3-opus", + custom_llm_provider="anthropic", + logging_obj=log_obj, + ) + + captured_ids = {} + + def fake_exception_type(**kwargs): + captured_ids["trace_id"] = trace_id_var.get() + captured_ids["session_id"] = session_id_var.get() + return ValueError("mapped boom") + + monkeypatch.setattr("litellm.litellm_core_utils.streaming_handler.exception_type", fake_exception_type) + + with pytest.raises(Exception): + wrapper._handle_stream_fallback_error(RuntimeError("boom")) + + # The mapper ran while the stream's own ids were still active. + assert captured_ids["trace_id"] == log_obj.litellm_trace_id + assert captured_ids["session_id"] == "fallback-error-session" + # Restored to the consumer's outer context once mapping/raise completes. + assert trace_id_var.get() == "outer-trace-fallback" + assert session_id_var.get() == "outer-session-fallback" + finally: + trace_id_var.set("") + session_id_var.set("") 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 6d318bb8729..7b7796fb01b 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -3562,6 +3562,8 @@ def test_supports_native_structured_outputs(): assert config._supports_native_structured_outputs("nvidia.nemotron-nano-3-30b") # DeepSeek: old substring "deepseek-v3.1" didn't match real ID assert config._supports_native_structured_outputs("deepseek.v3-v1:0") + assert config._supports_native_structured_outputs("deepseek.v3.2") + assert config._supports_native_structured_outputs("zai.glm-5") # Unsupported models -- should fall back to tool-call approach assert not config._supports_native_structured_outputs( diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 76bb11cc26d..2d1a5bc4c2a 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -2580,3 +2580,52 @@ def test_replayed_intercepted_search_turn_leaves_no_unsupported_block_for_bedroc assert "server_tool_use" not in serialized assert expected_evidence in serialized assert "Rome was founded in 753 BC." in serialized + + +@pytest.mark.parametrize("tool_type", ["web_search_20250305", "web_search_20260209"]) +def test_bedrock_invoke_messages_rejects_server_web_search_tool(tool_type: str): + """Bedrock can't execute Anthropic's server-side web search; the transform + must raise an actionable 400 pointing at the interception docs instead of + letting Bedrock return an opaque "provided request is not valid".""" + import litellm + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + with pytest.raises(litellm.BadRequestError) as exc_info: + cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "search the web for litellm"}], + anthropic_messages_optional_request_params={ + "max_tokens": 128, + "tools": [{"type": tool_type, "name": "web_search", "max_uses": 5}], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert "https://docs.litellm.ai/docs/integrations/websearch_interception" in str(exc_info.value) + assert "us.anthropic.claude-haiku-4-5-20251001-v1:0" in str(exc_info.value) + + +def test_bedrock_invoke_messages_allows_converted_websearch_function_tool(): + """The interception hook rewrites web_search into a plain custom tool + (litellm_web_search); that converted shape must pass through untouched.""" + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "search the web for litellm"}], + anthropic_messages_optional_request_params={ + "max_tokens": 128, + "tools": [ + { + "name": "litellm_web_search", + "description": "Search the web", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ], + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert result["tools"][0]["name"] == "litellm_web_search" diff --git a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py index 02fe7c8e68f..32141bead0e 100644 --- a/tests/test_litellm/llms/xai/test_xai_cost_calculator.py +++ b/tests/test_litellm/llms/xai/test_xai_cost_calculator.py @@ -432,10 +432,10 @@ class TestXAICostCalculator: model="grok-4.20-beta-0309-reasoning", usage=usage ) - # Input: 100 tokens * $2e-6 = $0.0002 - # Output: 200 tokens * $6e-6 = $0.0012 - expected_prompt_cost = 100 * 2e-6 - expected_completion_cost = 200 * 6e-6 + # Input: 100 tokens * $1.25e-6 = $0.000125 + # Output: 200 tokens * $2.5e-6 = $0.0005 + expected_prompt_cost = 100 * 1.25e-6 + expected_completion_cost = 200 * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) @@ -448,10 +448,38 @@ class TestXAICostCalculator: model="grok-4.20-beta-0309-non-reasoning", usage=usage ) - # Input: 50 tokens * $2e-6 = $0.0001 - # Output: 100 tokens * $6e-6 = $0.0006 - expected_prompt_cost = 50 * 2e-6 - expected_completion_cost = 100 * 6e-6 + # Input: 50 tokens * $1.25e-6 = $0.0000625 + # Output: 100 tokens * $2.5e-6 = $0.00025 + expected_prompt_cost = 50 * 1.25e-6 + expected_completion_cost = 100 * 2.5e-6 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + def test_grok_4_20_at_exactly_200k_prompt_tokens_uses_higher_tier(self): + """xAI bills the >=200k tier once the prompt reaches 200k, so the boundary is inclusive.""" + usage = Usage(prompt_tokens=200_000, completion_tokens=1_000, total_tokens=201_000) + + prompt_cost, completion_cost = cost_per_token( + model="grok-4.20-0309-reasoning", usage=usage + ) + + expected_prompt_cost = 200_000 * 2.5e-6 + expected_completion_cost = 1_000 * 5e-6 + + assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) + + def test_grok_4_20_just_below_200k_prompt_tokens_uses_base_tier(self): + """One token under the boundary still bills at the base rates.""" + usage = Usage(prompt_tokens=199_999, completion_tokens=1_000, total_tokens=200_999) + + prompt_cost, completion_cost = cost_per_token( + model="grok-4.20-0309-reasoning", usage=usage + ) + + expected_prompt_cost = 199_999 * 1.25e-6 + expected_completion_cost = 1_000 * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) @@ -464,10 +492,10 @@ class TestXAICostCalculator: model="grok-4.20-multi-agent-beta-0309", usage=usage ) - # Input: 200 tokens * $2e-6 = $0.0004 - # Output: 300 tokens * $6e-6 = $0.0018 - expected_prompt_cost = 200 * 2e-6 - expected_completion_cost = 300 * 6e-6 + # Input: 200 tokens * $1.25e-6 = $0.00025 + # Output: 300 tokens * $2.5e-6 = $0.00075 + expected_prompt_cost = 200 * 1.25e-6 + expected_completion_cost = 300 * 2.5e-6 assert math.isclose(prompt_cost, expected_prompt_cost, rel_tol=1e-10) assert math.isclose(completion_cost, expected_completion_cost, rel_tol=1e-10) diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 78b7e771239..9ff2c38d98a 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -3087,3 +3087,47 @@ class TestIsRequestBodySafeChecksBracketNotationMetadata: ) is True ) + + +class TestHasUserSetupSso: + """_has_user_setup_sso must treat SAML IdP metadata as SSO configured. + + Regression: UI discovery used this helper for sso_configured, but it only + checked OAuth client IDs, so SAML-only setups left the login button gray. + """ + + @pytest.fixture(autouse=True) + def _clear_sso_env(self, monkeypatch): + for key in ( + "MICROSOFT_CLIENT_ID", + "GOOGLE_CLIENT_ID", + "GENERIC_CLIENT_ID", + "SAML_IDP_METADATA_URL", + "SAML_IDP_METADATA_XML", + ): + monkeypatch.delenv(key, raising=False) + + def test_false_when_no_sso_env(self): + from litellm.proxy.auth.auth_utils import _has_user_setup_sso + + assert _has_user_setup_sso() is False + + def test_true_for_oauth_client_ids(self, monkeypatch): + from litellm.proxy.auth.auth_utils import _has_user_setup_sso + + monkeypatch.setenv("GOOGLE_CLIENT_ID", "google-client") + assert _has_user_setup_sso() is True + + def test_true_for_saml_metadata_url(self, monkeypatch): + from litellm.proxy.auth.auth_utils import _has_user_setup_sso + + monkeypatch.setenv( + "SAML_IDP_METADATA_URL", "https://idp.example.com/metadata.xml" + ) + assert _has_user_setup_sso() is True + + def test_true_for_saml_metadata_xml(self, monkeypatch): + from litellm.proxy.auth.auth_utils import _has_user_setup_sso + + monkeypatch.setenv("SAML_IDP_METADATA_XML", "") + assert _has_user_setup_sso() is True diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 616ad8a0981..608dc8cb5c8 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -2,7 +2,6 @@ import asyncio import json import os import sys -import time import types from datetime import datetime, timedelta, timezone from datetime import time as dt_time @@ -13,33 +12,19 @@ import pytest sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path -from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings -from litellm.proxy.utils import ProxyLogging # Mock classes for testing -class MockLiteLLMTeamMembership: - async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: - # Mock the update_many method for litellm_teammembership - return {"count": 1} +class MockTable: + """A single prisma table: records reads/writes and replays canned rows.""" - -class MockLiteLLMVerificationToken: def __init__(self): - self.update_many_calls: List[Dict[str, Any]] = [] - - async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: - self.update_many_calls.append({"where": where, "data": data}) - return {"count": 1} - - -class MockLiteLLMOrganizationTable: - def __init__(self): - self.update_many_calls: List[Dict[str, Any]] = [] self.find_many_calls: List[Dict[str, Any]] = [] + self.update_many_calls: List[Dict[str, Any]] = [] self._find_many_results: List[Any] = [] def set_find_many_results(self, results: List[Any]): @@ -54,43 +39,12 @@ class MockLiteLLMOrganizationTable: return {"count": 1} -class MockLiteLLMTagTable: - def __init__(self): - self.update_many_calls: List[Dict[str, Any]] = [] - self.find_many_calls: List[Dict[str, Any]] = [] - self._find_many_results: List[Any] = [] - - def set_find_many_results(self, results: List[Any]): - self._find_many_results = results - - async def find_many(self, where: Dict[str, Any]) -> List[Any]: - self.find_many_calls.append({"where": where}) - return self._find_many_results - - async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: - self.update_many_calls.append({"where": where, "data": data}) - return {"count": 1} - - -class MockLiteLLMEndUserTable: - def __init__(self): - self.find_many_calls: List[Dict[str, Any]] = [] - self._find_many_results: List[Any] = [] - - def set_find_many_results(self, results: List[Any]): - self._find_many_results = results - - async def find_many(self, where: Dict[str, Any]) -> List[Any]: - self.find_many_calls.append({"where": where}) - return self._find_many_results - - class MockBatcher: - """Captures per-row update calls and exposes them after commit(). + """Captures the writes queued on one `db.batch_()` and whether it committed. - Mirrors prisma's `db.batch_()` ergonomics enough that the reset job's - narrow-write helpers (`_write_key_reset_updates` et al) can run against - the mock and the test can assert on what would have been written. + Mirrors prisma's batch ergonomics enough that the reset job's write helpers + can run against the mock, and keeps `committed` so tests can prove a failed + cascade persisted nothing. """ def __init__(self): @@ -102,12 +56,23 @@ class MockBatcher: _self._table_name = table_name _self._outer = outer + def _record(_self, op, where, data): + _self._outer.calls.append({"table": _self._table_name, "op": op, "where": where, "data": data}) + def update(_self, where, data): - _self._outer.calls.append({"table": _self._table_name, "where": where, "data": data}) + _self._record("update", where, data) + + def update_many(_self, where, data): + _self._record("update_many", where, data) self.litellm_verificationtoken = _Table("key", self) self.litellm_usertable = _Table("user", self) self.litellm_teamtable = _Table("team", self) + self.litellm_budgettable = _Table("budget", self) + self.litellm_teammembership = _Table("team_membership", self) + self.litellm_organizationtable = _Table("org", self) + self.litellm_tagtable = _Table("tag", self) + self.litellm_endusertable = _Table("enduser", self) async def commit(self): self.committed = True @@ -116,16 +81,20 @@ class MockBatcher: class MockDB: def __init__(self): - self.litellm_teammembership = MockLiteLLMTeamMembership() - self.litellm_verificationtoken = MockLiteLLMVerificationToken() - self.litellm_endusertable = MockLiteLLMEndUserTable() - self.litellm_organizationtable = MockLiteLLMOrganizationTable() - self.litellm_tagtable = MockLiteLLMTagTable() + self.litellm_teammembership = MockTable() + self.litellm_verificationtoken = MockTable() + self.litellm_endusertable = MockTable() + self.litellm_organizationtable = MockTable() + self.litellm_tagtable = MockTable() self.batch_calls: List[Dict[str, Any]] = [] + self.batchers: List[MockBatcher] = [] def batch_(self): batcher = MockBatcher() - # Aggregate calls across all batches so tests can assert on cumulative writes. + self.batchers.append(batcher) + # Aggregate calls across all batches so tests can assert on cumulative + # writes. Only committed batches contribute: an abandoned batch writes + # nothing, exactly as prisma behaves. original_commit = batcher.commit async def _record_and_commit(): @@ -152,9 +121,11 @@ class MockPrismaClient: "budget": [], "enduser": [], } + self.get_data_calls: List[Dict[str, Any]] = [] self.db = MockDB() async def get_data(self, table_name, query_type, **kwargs): + self.get_data_calls.append({"table_name": table_name, "query_type": query_type, **kwargs}) data = self.data.get(table_name, []) # Handle specific filtering for budget table queries @@ -218,6 +189,39 @@ async def run_async_test(coro): return await coro +_ALREADY_EXPIRED = object() + + +def _budget_row( + budget_id: str = "test-budget-1", + budget_duration: Any = "7d", + budget_reset_at: Any = _ALREADY_EXPIRED, + max_budget: float = 10.0, +): + """An expiring budget tier, shaped like the rows get_data() hands back.""" + now = datetime.now(timezone.utc) + return type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": max_budget, + "budget_duration": budget_duration, + "budget_reset_at": (now - timedelta(hours=1) if budget_reset_at is _ALREADY_EXPIRED else budget_reset_at), + "budget_id": budget_id, + "created_at": now - timedelta(days=30), + }, + ) + + +def _batch_writes(mock_prisma_client, table: str, op: str | None = None) -> List[Dict[str, Any]]: + """Writes that were committed to the DB, optionally narrowed to one op.""" + return [ + call + for call in mock_prisma_client.db.batch_calls + if call["table"] == table and (op is None or call["op"] == op) + ] + + # Tests def test_write_key_reset_updates_skips_none_token_and_still_writes_the_rest(reset_budget_job, mock_prisma_client): """A key with token=None must be skipped, not queued as where={"token": None}. @@ -234,10 +238,10 @@ def test_write_key_reset_updates_skips_none_token_and_still_writes_the_rest(rese asyncio.run(reset_budget_job._write_key_reset_updates(updated_keys=keys)) - key_writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == "key"] - assert key_writes == [ + assert _batch_writes(mock_prisma_client, "key") == [ { "table": "key", + "op": "update", "where": {"token": "tok-ok"}, "data": {"spend": 0, "budget_reset_at": reset_at}, } @@ -369,18 +373,9 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): - # Setup test data + """End-user spend is zeroed and the tier's window advances, in one batch.""" now = datetime.now(timezone.utc) - test_budget = type( - "LiteLLM_BudgetTable", - (), - { - "max_budget": 500.0, - "budget_duration": "1d", - "budget_reset_at": now, - "budget_id": "test-budget-1", - }, - ) + test_budget = _budget_row(budget_id="test-budget-1", budget_duration="1d", budget_reset_at=now) test_enduser = type( "LiteLLM_EndUserTable", @@ -395,16 +390,22 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): mock_prisma_client.data["budget"] = [test_budget] mock_prisma_client.data["enduser"] = [test_enduser] - # Run the test asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Verify results - assert len(mock_prisma_client.updated_data["enduser"]) == 1 - assert len(mock_prisma_client.updated_data["budget"]) == 1 - updated_enduser = mock_prisma_client.updated_data["enduser"][0] - updated_budget = mock_prisma_client.updated_data["budget"][0] - assert updated_enduser.spend == 0.0 - assert updated_budget.budget_reset_at > now + assert _batch_writes(mock_prisma_client, "enduser") == [ + { + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["test-enduser-1"]}}, + "data": {"spend": 0}, + } + ] + + budget_writes = _batch_writes(mock_prisma_client, "budget") + assert len(budget_writes) == 1 + assert budget_writes[0]["where"] == {"budget_id": "test-budget-1"} + assert budget_writes[0]["data"]["budget_reset_at"] > now + assert set(budget_writes[0]["data"].keys()) == {"budget_reset_at"} def test_reset_budget_all(reset_budget_job, mock_prisma_client): @@ -485,190 +486,81 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): ("user", {"user_id": "uid-all-1"}), ("team", {"team_id": "tid-all-1"}), ]: - writes = [c for c in mock_prisma_client.db.batch_calls if c["table"] == table_name] + writes = _batch_writes(mock_prisma_client, table_name, op="update") assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}" assert writes[0]["where"] == where assert writes[0]["data"]["spend"] == 0 assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"} - # Enduser + budget rows still go through update_data (not narrowed; different path). - assert len(mock_prisma_client.updated_data["enduser"]) == 1 - assert len(mock_prisma_client.updated_data["budget"]) == 1 - assert mock_prisma_client.updated_data["enduser"][0].spend == 0.0 - - -def test_reset_budget_for_keys_linked_to_budgets(reset_budget_job, mock_prisma_client): - """ - Test that when a budget tier is reset, keys linked to that budget - (via budget_id) that don't have their own budget_duration also get - their spend reset. - - This covers the case where keys were created with budget_id but - budget_duration was not inherited to the key (pre-fix keys). - """ - from litellm.proxy._types import LiteLLM_BudgetTableFull - - now = datetime.now(timezone.utc) - - # Create a budget tier that is due for reset - test_budget = type( - "LiteLLM_BudgetTableFull", - (), + # The budget tier's cascade rides the same batch machinery. + assert _batch_writes(mock_prisma_client, "enduser") == [ { - "max_budget": 10.0, - "budget_duration": "7d", - "budget_reset_at": now - timedelta(hours=1), - "budget_id": "7d-budget-tier", - "created_at": now - timedelta(days=7), + "table": "enduser", + "op": "update_many", + "where": {"user_id": {"in": ["test-enduser-1"]}}, + "data": {"spend": 0}, + } + ] + assert len(_batch_writes(mock_prisma_client, "budget")) == 1 + + +_LINKED_TABLE_CASES = [ + ("team_membership", {"budget_id": {"in": ["7d-budget-tier"]}}), + ( + "key", + { + "budget_id": {"in": ["7d-budget-tier"]}, + "budget_duration": None, + "spend": {"gt": 0}, }, - ) - - budgets_to_reset = [test_budget] - - # Run the method - asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset)) - - # Verify that update_many was called on litellm_verificationtoken - calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls - assert len(calls) == 1, f"Expected 1 update_many call, got {len(calls)}" - - # Verify the where clause filters by budget_id and null budget_duration - call = calls[0] - assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]} - assert call["where"]["budget_duration"] is None - - # Verify spend is reset to 0 - assert call["data"]["spend"] == 0 + ), + ("org", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), + ("tag", {"budget_id": {"in": ["7d-budget-tier"]}, "spend": {"gt": 0}}), +] -def test_reset_budget_for_keys_linked_to_budgets_excludes_keys_with_own_budget_duration( - reset_budget_job, mock_prisma_client +@pytest.mark.parametrize( + "table, expected_where", + _LINKED_TABLE_CASES, + ids=[case[0] for case in _LINKED_TABLE_CASES], +) +def test_budget_table_reset_zeroes_spend_on_every_linked_table( + reset_budget_job, mock_prisma_client, table, expected_where ): + """One expiring tier zeroes spend on every row it gates. + + The filters carry real behavior: keys must be narrowed to + `budget_duration: None` so keys with their own reset schedule aren't + double-reset by reset_budget_for_litellm_keys(), and the payload must stay + exactly {"spend": 0} because `total_spend` is a lifetime counter a reset + may never touch. """ - Test that keys with BOTH budget_id AND budget_duration are excluded from - reset_budget_for_keys_linked_to_budgets. Such keys have their own reset - schedule and are handled only by reset_budget_for_litellm_keys(). The - budget_duration=None filter ensures they are NOT double-reset when the - linked budget tier expires. - """ - now = datetime.now(timezone.utc) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="7d-budget-tier", budget_duration="7d")] - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "max_budget": 10.0, - "budget_duration": "7d", - "budget_reset_at": now - timedelta(hours=1), - "budget_id": "7d-budget-tier", - "created_at": now - timedelta(days=7), - }, - ) + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - budgets_to_reset = [test_budget] - - asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset)) - - calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls - assert len(calls) == 1 - call = calls[0] - - # Critical: budget_duration must be None so keys with their own budget_duration - # (e.g. key has budget_id="X" AND budget_duration=60) are excluded. - # Those keys are reset only by reset_budget_for_litellm_keys() - no double-reset. - assert call["where"]["budget_duration"] is None - assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]} + writes = _batch_writes(mock_prisma_client, table, op="update_many") + assert len(writes) == 1, f"expected exactly 1 {table} write, got {writes}" + assert writes[0]["where"] == expected_where + assert writes[0]["data"] == {"spend": 0} -def test_reset_budget_for_keys_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): - """ - Test that when there are no budgets to reset, no update is performed - on the verification token table. - """ - # Run with empty list - asyncio.run(reset_budget_job.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=[])) +def test_budget_table_reset_writes_nothing_when_no_budget_is_due(reset_budget_job, mock_prisma_client): + """Nothing due means no transaction is opened at all.""" + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Verify no update_many calls were made - calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls - assert len(calls) == 0 + assert mock_prisma_client.db.batchers == [] + assert mock_prisma_client.db.batch_calls == [] -def test_reset_budget_for_orgs_linked_to_budgets(reset_budget_job, mock_prisma_client): - """ - Test that when a budget tier is reset, orgs linked to that budget - (via budget_id) also get their spend reset. - """ - now = datetime.now(timezone.utc) +def _run_reset_at_fixed_now(job, fixed_now): + """Run the budget-table reset with `now` pinned for reset-time math.""" + from unittest.mock import patch - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "max_budget": 100.0, - "budget_duration": "30d", - "budget_reset_at": now - timedelta(hours=1), - "budget_id": "30d-org-budget", - "created_at": now - timedelta(days=30), - }, - ) - - asyncio.run(reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[test_budget])) - - calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls - assert len(calls) == 1 - call = calls[0] - assert call["where"]["budget_id"] == {"in": ["30d-org-budget"]} - assert call["where"]["spend"] == {"gt": 0} - assert call["data"]["spend"] == 0 - - -def test_reset_budget_for_orgs_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): - """ - Test that when there are no budgets to reset, no update is performed - on the organization table. - """ - asyncio.run(reset_budget_job.reset_budget_for_orgs_linked_to_budgets(budgets_to_reset=[])) - calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls - assert len(calls) == 0 - - -def test_reset_budget_for_tags_linked_to_budgets(reset_budget_job, mock_prisma_client): - """ - Test that when a budget tier is reset, tags linked to that budget - (via budget_id) also get their spend reset. - """ - now = datetime.now(timezone.utc) - - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "max_budget": 50.0, - "budget_duration": "30d", - "budget_reset_at": now - timedelta(hours=1), - "budget_id": "30d-tag-budget", - "created_at": now - timedelta(days=30), - }, - ) - - asyncio.run(reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[test_budget])) - - calls = mock_prisma_client.db.litellm_tagtable.update_many_calls - assert len(calls) == 1 - call = calls[0] - assert call["where"]["budget_id"] == {"in": ["30d-tag-budget"]} - assert call["where"]["spend"] == {"gt": 0} - assert call["data"]["spend"] == 0 - - -def test_reset_budget_for_tags_linked_to_budgets_empty(reset_budget_job, mock_prisma_client): - """ - Test that when there are no budgets to reset, no update is performed - on the tag table. - """ - asyncio.run(reset_budget_job.reset_budget_for_tags_linked_to_budgets(budgets_to_reset=[])) - calls = mock_prisma_client.db.litellm_tagtable.update_many_calls - assert len(calls) == 0 + with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: + mock_dt.now.return_value = fixed_now + mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) + asyncio.run(job.reset_budget_for_litellm_budget_table()) @pytest.mark.parametrize( @@ -680,215 +572,70 @@ def test_reset_budget_for_tags_linked_to_budgets_empty(reset_budget_job, mock_pr ], ids=["30d-calendar-month", "1mo-calendar-month", "1d-next-midnight"], ) -def test_reset_budget_reset_at_date_calendar_aligned(budget_duration, expected_day, expected_month): - """ - Verify that _reset_budget_reset_at_date produces calendar-aligned reset - times (matching get_budget_reset_time), not sliding-window offsets. - """ - from unittest.mock import patch - - # Fix "now" to 2023-06-15 10:30:00 UTC for deterministic results +def test_budget_reset_at_written_is_calendar_aligned( + reset_budget_job, mock_prisma_client, budget_duration, expected_day, expected_month +): + """The advanced budget_reset_at is calendar-aligned, not a sliding + now + duration offset.""" fixed_now = datetime(2023, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + mock_prisma_client.data["budget"] = [ + _budget_row( + budget_id="test-budget", + budget_duration=budget_duration, + budget_reset_at=fixed_now - timedelta(hours=1), + ) + ] - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "budget_duration": budget_duration, - "budget_reset_at": fixed_now - timedelta(hours=1), - "budget_id": "test-budget", - "created_at": fixed_now - timedelta(days=30), - }, - ) + _run_reset_at_fixed_now(reset_budget_job, fixed_now) - with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: - mock_dt.now.return_value = fixed_now - mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) - - assert test_budget.budget_reset_at.day == expected_day - assert test_budget.budget_reset_at.month == expected_month - assert test_budget.budget_reset_at.hour == 0 - assert test_budget.budget_reset_at.minute == 0 - assert test_budget.budget_reset_at.second == 0 + writes = _batch_writes(mock_prisma_client, "budget") + assert len(writes) == 1 + written = writes[0]["data"]["budget_reset_at"] + assert (written.day, written.month) == (expected_day, expected_month) + assert (written.hour, written.minute, written.second) == (0, 0, 0) -def test_reset_budget_reset_at_date_7d_next_monday(): - """Verify 7d budget duration resets to next Monday at midnight.""" - from unittest.mock import patch - +def test_budget_reset_at_written_for_7d_is_next_monday(reset_budget_job, mock_prisma_client): + """7d budgets advance to next Monday at midnight.""" # 2023-06-14 is a Wednesday fixed_now = datetime(2023, 6, 14, 10, 30, 0, tzinfo=timezone.utc) + mock_prisma_client.data["budget"] = [ + _budget_row(budget_id="test-budget", budget_duration="7d", budget_reset_at=fixed_now - timedelta(hours=1)) + ] - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "budget_duration": "7d", - "budget_reset_at": fixed_now - timedelta(hours=1), - "budget_id": "test-budget", - "created_at": fixed_now - timedelta(days=7), - }, - ) + _run_reset_at_fixed_now(reset_budget_job, fixed_now) - with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: - mock_dt.now.return_value = fixed_now - mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) - - # Next Monday after Wednesday June 14 is June 19 - assert test_budget.budget_reset_at.day == 19 - assert test_budget.budget_reset_at.month == 6 - assert test_budget.budget_reset_at.weekday() == 0 # Monday - assert test_budget.budget_reset_at.hour == 0 + written = _batch_writes(mock_prisma_client, "budget")[0]["data"]["budget_reset_at"] + assert (written.day, written.month) == (19, 6) + assert written.weekday() == 0 + assert written.hour == 0 -def test_reset_budget_reset_at_date_none_duration(): - """Verify that budget_reset_at is unchanged when budget_duration is None.""" - original_reset_at = datetime(2023, 6, 20, 0, 0, 0, tzinfo=timezone.utc) - now = datetime(2023, 6, 15, 10, 0, 0, tzinfo=timezone.utc) +def test_budget_with_no_duration_gets_no_reset_at_write(reset_budget_job, mock_prisma_client): + """A tier without a duration has no next window, so its row is left alone + rather than rewritten with an unchanged value.""" + mock_prisma_client.data["budget"] = [ + _budget_row( + budget_id="no-duration", budget_duration=None, budget_reset_at=datetime(2023, 6, 20, tzinfo=timezone.utc) + ) + ] - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "budget_duration": None, - "budget_reset_at": original_reset_at, - "budget_id": "test-budget", - "created_at": now - timedelta(days=30), - }, - ) + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, now, BudgetResetSettings())) - assert test_budget.budget_reset_at == original_reset_at + assert _batch_writes(mock_prisma_client, "budget") == [] -def test_reset_budget_reset_at_date_none_reset_at(): - """Verify that budget_reset_at is set correctly even when previously None.""" - from unittest.mock import patch - +def test_budget_reset_at_written_when_previously_null(reset_budget_job, mock_prisma_client): + """A tier whose budget_reset_at was never initialized still gets one.""" fixed_now = datetime(2023, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + mock_prisma_client.data["budget"] = [ + _budget_row(budget_id="test-budget", budget_duration="30d", budget_reset_at=None) + ] - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "budget_duration": "30d", - "budget_reset_at": None, - "budget_id": "test-budget", - "created_at": fixed_now - timedelta(days=5), - }, - ) + _run_reset_at_fixed_now(reset_budget_job, fixed_now) - with patch("litellm.proxy.common_utils.timezone_utils.datetime") as mock_dt: - mock_dt.now.return_value = fixed_now - mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) - asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now, BudgetResetSettings())) - - # Should be set to 1st of next month (July 1) - assert test_budget.budget_reset_at is not None - assert test_budget.budget_reset_at.day == 1 - assert test_budget.budget_reset_at.month == 7 - - -def test_budget_table_reset_also_resets_linked_keys(reset_budget_job, mock_prisma_client): - """ - Integration-style test: when reset_budget_for_litellm_budget_table runs, - it should also reset spend for keys linked to the expiring budget tiers - (in addition to end-users and team members). - """ - now = datetime.now(timezone.utc) - - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "max_budget": 10.0, - "budget_duration": "7d", - "budget_reset_at": now - timedelta(hours=1), - "budget_id": "7d-budget-tier", - "created_at": now - timedelta(days=7), - }, - ) - - mock_prisma_client.data["budget"] = [test_budget] - - # Run the full budget table reset - asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - - # Verify that keys linked to the budget were also reset - calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls - assert len(calls) == 1, ( - "Expected reset_budget_for_litellm_budget_table to also reset keys " - f"linked to expiring budgets, but got {len(calls)} update_many calls" - ) - assert calls[0]["where"]["budget_id"] == {"in": ["7d-budget-tier"]} - assert calls[0]["data"]["spend"] == 0 - - -def test_budget_table_reset_also_resets_linked_orgs(reset_budget_job, mock_prisma_client): - """ - Integration-style test: when reset_budget_for_litellm_budget_table runs, - it should also reset spend for orgs linked to the expiring budget tiers - (in addition to end-users, team members, and keys). - """ - now = datetime.now(timezone.utc) - - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "max_budget": 100.0, - "budget_duration": "30d", - "budget_reset_at": now - timedelta(hours=1), - "budget_id": "30d-org-budget", - "created_at": now - timedelta(days=30), - }, - ) - - mock_prisma_client.data["budget"] = [test_budget] - - asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - - calls = mock_prisma_client.db.litellm_organizationtable.update_many_calls - assert len(calls) == 1, ( - "Expected reset_budget_for_litellm_budget_table to also reset orgs " - f"linked to expiring budgets, but got {len(calls)} update_many calls" - ) - assert calls[0]["where"]["budget_id"] == {"in": ["30d-org-budget"]} - assert calls[0]["data"]["spend"] == 0 - - -def test_budget_table_reset_also_resets_linked_tags(reset_budget_job, mock_prisma_client): - """ - Integration-style test: when reset_budget_for_litellm_budget_table runs, - it should also reset spend for tags linked to the expiring budget tiers. - """ - now = datetime.now(timezone.utc) - - test_budget = type( - "LiteLLM_BudgetTableFull", - (), - { - "max_budget": 50.0, - "budget_duration": "30d", - "budget_reset_at": now - timedelta(hours=1), - "budget_id": "30d-tag-budget", - "created_at": now - timedelta(days=30), - }, - ) - - mock_prisma_client.data["budget"] = [test_budget] - - asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - - calls = mock_prisma_client.db.litellm_tagtable.update_many_calls - assert len(calls) == 1, ( - "Expected reset_budget_for_litellm_budget_table to also reset tags " - f"linked to expiring budgets, but got {len(calls)} update_many calls" - ) - assert calls[0]["where"]["budget_id"] == {"in": ["30d-tag-budget"]} - assert calls[0]["data"]["spend"] == 0 + written = _batch_writes(mock_prisma_client, "budget")[0]["data"]["budget_reset_at"] + assert (written.day, written.month) == (1, 7) def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock_prisma_client): @@ -965,16 +712,14 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Both end users should have been reset - updated = mock_prisma_client.updated_data["enduser"] - assert len(updated) == 2, f"Expected 2 endusers reset (1 explicit + 1 implicit), got {len(updated)}" - - user_ids = {u.user_id for u in updated} - assert "enduser-explicit" in user_ids - assert "enduser-implicit" in user_ids - - for u in updated: - assert u.spend == 0.0, f"Expected spend=0 for {u.user_id}, got {u.spend}" + # Both end users are zeroed by the same committed statement. + enduser_writes = _batch_writes(mock_prisma_client, "enduser") + assert len(enduser_writes) == 1, f"Expected a single enduser write, got {enduser_writes}" + assert set(enduser_writes[0]["where"]["user_id"]["in"]) == { + "enduser-explicit", + "enduser-implicit", + } + assert enduser_writes[0]["data"] == {"spend": 0} # Verify find_many was called to fetch NULL-budget-id end users find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls @@ -1054,34 +799,6 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li litellm.max_end_user_budget_id = None -def test_reset_budget_for_team_members_preserves_total_spend(): - """Regression guard: reset_budget_for_litellm_team_members must zero `spend` - but leave `total_spend` untouched. - - The reset writes `data={"spend": 0}` explicitly. If a future refactor adds - `"total_spend": 0` to that dict, this test fails immediately. - """ - expired_budget = type( - "LiteLLM_BudgetTableFull", - (), - {"budget_id": "budget-1"}, - ) - - mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=mock_prisma_client) - - asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) - - mock_prisma_client.db.litellm_teammembership.update_many.assert_called_once() - call_kwargs = mock_prisma_client.db.litellm_teammembership.update_many.call_args.kwargs - assert call_kwargs["where"]["budget_id"]["in"] == ["budget-1"] - assert call_kwargs["data"] == {"spend": 0} - assert "total_spend" not in call_kwargs["data"] - - # --------------------------------------------------------------------------- # reset_budget_windows (per-key / per-team concurrent window resets) # --------------------------------------------------------------------------- @@ -1323,28 +1040,6 @@ def _make_counter_invalidation_job(monkeypatch): return spend_counter_cache -def test_reset_budget_for_team_members_invalidates_redis_counter(monkeypatch): - """Team-member budget reset clears the Redis spend counter.""" - counter_cache = _make_counter_invalidation_job(monkeypatch) - - expired_budget = type("B", (), {"budget_id": "budget-1"}) - membership = type( - "Membership", - (), - {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"}, - ) - - prisma_client = MagicMock() - prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[membership]) - prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) - - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team_member:alice:team-x", value=0.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:team_member:alice:team-x", value=0.0, ttl=60) - - def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): """Key budget reset must clear the Redis spend counter.""" counter_cache = _make_counter_invalidation_job(monkeypatch) @@ -1574,207 +1269,240 @@ def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job, ) -def test_reset_budget_for_keys_linked_to_budgets_invalidates_redis_counter(monkeypatch): - """Resetting keys via budget tier must clear each linked key's counter.""" - counter_cache = _make_counter_invalidation_job(monkeypatch) - - expired_budget = type("B", (), {"budget_id": "budget-1"}) - linked_key = type("Key", (), {"token": "sk-linked"}) - - prisma_client = MagicMock() - prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[linked_key]) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(return_value={"count": 1}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) - - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-linked", value=0.0, ttl=60) +_INVALIDATION_CASES = [ + ( + "litellm_teammembership", + type("Membership", (), {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"}), + "spend:team_member:alice:team-x", + {"team-x_alice"}, + ), + ( + "litellm_verificationtoken", + type("Key", (), {"token": "sk-linked"}), + "spend:key:sk-linked", + {"sk-linked"}, + ), + ( + "litellm_organizationtable", + type("Org", (), {"organization_id": "org-acme"}), + "spend:org:org-acme", + {"org_id:org-acme", "org_id:org-acme:with_budget"}, + ), + ( + "litellm_tagtable", + type("Tag", (), {"tag_name": "tenant-42"}), + "spend:tag:tenant-42", + {"tag:tenant-42"}, + ), +] -def test_reset_budget_for_orgs_linked_to_budgets_invalidates_redis_counter(monkeypatch): - """Resetting orgs via budget tier must clear each linked org's counter.""" - counter_cache = _make_counter_invalidation_job(monkeypatch) - - expired_budget = type("B", (), {"budget_id": "budget-1"}) - linked_org = type("Org", (), {"organization_id": "org-acme"}) - - prisma_client = MagicMock() - prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[linked_org]) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock(return_value={"count": 1}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget])) - - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:org:org-acme", value=0.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:org:org-acme", value=0.0, ttl=60) - - -def test_reset_budget_for_tags_linked_to_budgets_invalidates_redis_counter(monkeypatch): - """Resetting tags via budget tier must clear each linked tag's counter.""" - counter_cache = _make_counter_invalidation_job(monkeypatch) - - expired_budget = type("B", (), {"budget_id": "budget-1"}) - linked_tag = type("Tag", (), {"tag_name": "tenant-42"}) - - prisma_client = MagicMock() - prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag]) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) - - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:tag:tenant-42", value=0.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:tag:tenant-42", value=0.0, ttl=60) - - -def test_reset_budget_for_tags_linked_to_budgets_invalidates_management_cache( - monkeypatch, +@pytest.mark.parametrize( + "table_attr, linked_row, counter_key, cache_keys", + _INVALIDATION_CASES, + ids=["team_membership", "key", "org", "tag"], +) +def test_budget_table_reset_invalidates_counters_and_management_cache( + reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys ): - """Regression guard for the bug where tag spend stayed frozen across cycles. + """Every row the cascade zeroes gets its spend counter cleared and its + management-cache entry dropped. - ``SpendCounterReseed.from_db`` returns ``None`` for ``spend:tag:*`` keys, - so once the spend counter expires the tag budget check falls back to the - cached ``LiteLLM_TagTable.spend``. If we don't drop the management cache - entry on reset, that cached object lingers (TTL 60s) with the pre-reset - spend, and ``_tag_max_budget_check`` keeps returning HTTP 400 even though - the DB row has been zeroed. + Both matter. ``SpendCounterReseed.from_db`` returns None for tags, so once + the counter expires the budget check falls back to the cached row's + ``.spend``; and for keys, orgs and team memberships another pod's cached + object can stay pinned above the zeroed DB row until its TTL. Team + membership cache keys follow auth's ``{team_id}_{user_id}`` shape, and orgs + carry both the plain and the ``:with_budget`` entry. """ counter_cache = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + getattr(mock_prisma_client.db, table_attr).set_find_many_results([linked_row]) - expired_budget = type("B", (), {"budget_id": "budget-1"}) - linked_tag = type("Tag", (), {"tag_name": "tenant-42"}) + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - prisma_client = MagicMock() - prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag]) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) - - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="tag:tenant-42") + counter_cache.in_memory_cache.set_cache.assert_any_call(key=counter_key, value=0.0, ttl=60) + counter_cache.redis_cache.async_set_cache.assert_any_await(key=counter_key, value=0.0, ttl=60) + deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + assert cache_keys <= deleted -def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management_cache( - monkeypatch, -): - """When multiple tags share the expired budget tier, every one of them - has its ``user_api_key_cache`` entry dropped — not just the first.""" +def test_budget_table_reset_invalidates_every_tag_not_just_the_first(reset_budget_job, mock_prisma_client, monkeypatch): + """When several tags share the expiring tier, all of them are evicted.""" counter_cache = _make_counter_invalidation_job(monkeypatch) - - expired_budget = type("B", (), {"budget_id": "budget-1"}) - linked_tags = [ - type("Tag", (), {"tag_name": "tenant-a"}), - type("Tag", (), {"tag_name": "tenant-b"}), - type("Tag", (), {"tag_name": "tenant-c"}), - ] - - prisma_client = MagicMock() - prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=linked_tags) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 3}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) - - deleted_keys = { - call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list - } - assert deleted_keys == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"} - - -def test_reset_budget_for_keys_linked_to_budgets_invalidates_management_cache( - monkeypatch, -): - """Budget-tier key resets must drop the cached key object (hashed token key). - - Historically this test used ``assert_not_awaited()`` on - ``user_api_key_cache.async_delete_cache``, reflecting the assumption that - ``SpendCounterReseed.from_db`` alone kept spend consistent for keys and - that invalidating the management cache was unnecessary. That was flipped to - ``assert_any_await(...)`` because the old invariant fails across pods: a - budget reset on one instance can leave another pod's cached key object - (including embedded ``.spend``) stale until TTL expiry. Eviction now matches - tags/orgs/teams. Do not treat the ``cache_key_fn`` / invalidation wiring as - redundant without revisiting that cross-pod consistency story. - """ - counter_cache = _make_counter_invalidation_job(monkeypatch) - - expired_budget = type("B", (), {"budget_id": "budget-1"}) - linked_key = type("Key", (), {"token": "sk-linked"}) - - prisma_client = MagicMock() - prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[linked_key]) - prisma_client.db.litellm_verificationtoken.update_many = AsyncMock(return_value={"count": 1}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget])) - - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="sk-linked") - - -def test_reset_budget_for_orgs_linked_to_budgets_invalidates_management_cache( - monkeypatch, -): - """Org rows use both base and budget-table cache keys — evict both on reset.""" - counter_cache = _make_counter_invalidation_job(monkeypatch) - - expired_budget = type("B", (), {"budget_id": "budget-1"}) - linked_org = type("Org", (), {"organization_id": "org-acme"}) - - prisma_client = MagicMock() - prisma_client.db.litellm_organizationtable.find_many = AsyncMock(return_value=[linked_org]) - prisma_client.db.litellm_organizationtable.update_many = AsyncMock(return_value={"count": 1}) - - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget])) - - deleted_keys = { - call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list - } - assert deleted_keys == { - "org_id:org-acme", - "org_id:org-acme:with_budget", - } - - -def test_reset_budget_for_team_members_invalidates_management_cache(monkeypatch): - """Team membership cache key matches auth: ``{team_id}_{user_id}``.""" - counter_cache = _make_counter_invalidation_job(monkeypatch) - - expired_budget = type("B", (), {"budget_id": "budget-1"}) - membership = type( - "Membership", - (), - {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"}, + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + mock_prisma_client.db.litellm_tagtable.set_find_many_results( + [type("Tag", (), {"tag_name": name}) for name in ("tenant-a", "tenant-b", "tenant-c")] ) - prisma_client = MagicMock() - prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[membership]) - prisma_client.db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 1}) + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget])) - - counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(key="team-x_alice") + deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} + assert deleted == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"} -def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure_still_resets( - monkeypatch, -): - """If ``async_delete_cache`` raises, the DB cascade must still complete.""" +def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_job, mock_prisma_client, monkeypatch): + """Eviction runs after the commit, so a broken cache cannot undo the write.""" counter_cache = _make_counter_invalidation_job(monkeypatch) counter_cache.user_api_key_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("cache unavailable")) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + mock_prisma_client.db.litellm_tagtable.set_find_many_results([type("Tag", (), {"tag_name": "tenant-42"})]) - expired_budget = type("B", (), {"budget_id": "budget-1"}) - linked_tag = type("Tag", (), {"tag_name": "tenant-42"}) + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - prisma_client = MagicMock() - prisma_client.db.litellm_tagtable.find_many = AsyncMock(return_value=[linked_tag]) - prisma_client.db.litellm_tagtable.update_many = AsyncMock(return_value={"count": 1}) + assert len(_batch_writes(mock_prisma_client, "tag", op="update_many")) == 1 + assert mock_prisma_client.db.batchers[0].committed is True - job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client) - asyncio.run(job.reset_budget_for_tags_linked_to_budgets([expired_budget])) - prisma_client.db.litellm_tagtable.update_many.assert_awaited_once() +# --------------------------------------------------------------------------- +# Atomicity of the budget-table cascade (LIT-5138) +# --------------------------------------------------------------------------- + + +class FailingCommitDB(MockDB): + """Batches that blow up at commit, like a Postgres timeout mid-cascade.""" + + def batch_(self): + batcher = super().batch_() + + async def _fail(): + raise RuntimeError("simulated Postgres timeout mid-cascade") + + batcher.commit = _fail + return batcher + + +class FailingTeamMembershipDB(MockDB): + """Queueing the team-membership reset raises, i.e. the cascade breaks after + earlier writes are already queued.""" + + def batch_(self): + batcher = super().batch_() + + def _fail(where, data): + raise RuntimeError("simulated failure queueing the team-membership reset") + + batcher.litellm_teammembership.update_many = _fail + return batcher + + +class OrderRecordingDB(MockDB): + """Appends a marker to a shared list when a batch commits.""" + + def __init__(self, events): + super().__init__() + self._events = events + + def batch_(self): + batcher = super().batch_() + wrapped = batcher.commit + + async def _record_commit(): + self._events.append("commit") + return await wrapped() + + batcher.commit = _record_commit + return batcher + + +def _job_with_expired_budget(db, proxy_logging=None): + """A job with one due tier and a linked tag, so cache invalidation has + something to invalidate and its absence is a real signal.""" + prisma_client = MockPrismaClient() + prisma_client.db = db + prisma_client.data["budget"] = [_budget_row(budget_id="budget-1", budget_duration="7d")] + db.litellm_tagtable.set_find_many_results([type("Tag", (), {"tag_name": "tenant-42"})]) + job = ResetBudgetJob( + proxy_logging_obj=proxy_logging or MockProxyLogging(), + prisma_client=prisma_client, + ) + return job, prisma_client + + +@pytest.mark.parametrize( + "db_factory", + [FailingCommitDB, FailingTeamMembershipDB], + ids=["commit-fails", "queueing-fails"], +) +def test_budget_reset_at_is_not_advanced_when_the_cascade_fails(db_factory, monkeypatch): + """Regression for LIT-5138. + + The old code committed the new budget_reset_at first and zeroed the + dependent spend afterwards. A failure part-way through left the tier + stamped for the next window, so every later tick skipped it and team + member / enduser / org / tag spend stayed at the cap for the whole window. + One transaction means a failure anywhere persists nothing and the tier is + still due on the next tick. + """ + counter_cache = _make_counter_invalidation_job(monkeypatch) + job, prisma_client = _job_with_expired_budget(db_factory()) + + asyncio.run(job.reset_budget_for_litellm_budget_table()) # swallowed, retried next tick + + assert prisma_client.db.batch_calls == [], "a failed cascade must not persist any write" + assert prisma_client.db.batchers[0].committed is False + assert prisma_client.updated_data["budget"] == [], "budget_reset_at must not be advanced outside the transaction" + counter_cache.in_memory_cache.set_cache.assert_not_called() + counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited() + + +def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mock_prisma_client, monkeypatch): + """Dependent spend and the budget_reset_at advance ride one batch.""" + _make_counter_invalidation_job(monkeypatch) + now = datetime.now(timezone.utc) + budget = _budget_row(budget_id="budget-1", budget_duration="7d") + mock_prisma_client.data["budget"] = [budget] + mock_prisma_client.data["enduser"] = [ + type("EndUser", (), {"spend": 5.0, "litellm_budget_table": budget, "user_id": "enduser-1"}) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + assert len(mock_prisma_client.db.batchers) == 1, "the cascade must not be split across transactions" + batcher = mock_prisma_client.db.batchers[0] + assert batcher.committed is True + assert {(call["table"], call["op"]) for call in batcher.calls} == { + ("team_membership", "update_many"), + ("key", "update_many"), + ("org", "update_many"), + ("tag", "update_many"), + ("enduser", "update_many"), + ("budget", "update_many"), + } + budget_write = next(call for call in batcher.calls if call["table"] == "budget") + assert budget_write["data"]["budget_reset_at"] > now + + +def test_caches_are_invalidated_only_after_the_transaction_commits(monkeypatch): + """A counter zeroed before the write lands would admit requests past the + cap while the DB still holds the over-budget spend.""" + events = [] + counter_cache = _make_counter_invalidation_job(monkeypatch) + counter_cache.in_memory_cache.set_cache.side_effect = lambda **kwargs: events.append("counter") + + job, _ = _job_with_expired_budget(OrderRecordingDB(events)) + + asyncio.run(job.reset_budget_for_litellm_budget_table()) + + assert events == ["commit", "counter"] + + +def test_failed_cascade_is_logged_as_a_cascade_failure(monkeypatch): + """The failure log has to name what actually broke. The old catch-all + blamed end users even when the team-membership write was the failure.""" + from unittest.mock import patch + + _make_counter_invalidation_job(monkeypatch) + job, _ = _job_with_expired_budget(FailingTeamMembershipDB()) + + with patch("litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception") as mock_exception: + asyncio.run(job.reset_budget_for_litellm_budget_table()) + + assert mock_exception.call_count == 1 + message = mock_exception.call_args.args[0] + assert "cascade" in message + for mentioned in ("team member", "enduser", "org", "tag", "budget_reset_at"): + assert mentioned in message, f"failure log should mention {mentioned}: {message}" def _extract_reset_where(find_many_mock): @@ -1799,23 +1527,65 @@ def _asserts_null_reset_is_due(where): branches = where.get("OR") assert isinstance(branches, list), f"expected an OR filter, got {where!r}" - has_null_branch = any( - b.get("AND") - == [ - {"budget_reset_at": None}, - {"NOT": {"budget_duration": None}}, - ] - for b in branches - if isinstance(b, dict) - ) - has_expired_branch = any( - isinstance(b, dict) - and "budget_reset_at" in b - and b["budget_reset_at"] is not None - for b in branches - ) + has_null_branch = {"budget_reset_at": None} in branches + has_expired_branch = any(isinstance(b, dict) and isinstance(b.get("budget_reset_at"), dict) for b in branches) assert has_null_branch, f"missing NULL-reset_at branch in {where!r}" assert has_expired_branch, f"missing expired-reset_at branch in {where!r}" + assert where.get("NOT") == {"budget_duration": None}, f"NULL reset_at is only due with a duration: {where!r}" + + +_RESET_TABLE_ATTRS = { + "user": "litellm_usertable", + "team": "litellm_teamtable", + "budget": "litellm_budgettable", + "key": "litellm_verificationtoken", +} + + +def _run_reset_query(table_name, **extra): + """Run ``get_data`` for one table's budget-reset query against a mocked + prisma handle, and hand back the ``find_many`` mock it drove.""" + from litellm.proxy.utils import PrismaClient + + client = PrismaClient.__new__(PrismaClient) + client.db = MagicMock() + find_many = AsyncMock(return_value=[]) + setattr(getattr(client.db, _RESET_TABLE_ATTRS[table_name]), "find_many", find_many) + + now = datetime.now(timezone.utc) + expires = {"expires": now} if table_name == "key" else {} + asyncio.run(client.get_data(table_name=table_name, query_type="find_all", reset_at=now, **expires, **extra)) + return find_many + + +@pytest.mark.parametrize("table_name", ["user", "team", "budget", "key"]) +def test_get_data_reset_query_applies_the_row_limit(table_name): + """The reset job pages through due rows, so ``limit`` has to reach prisma as + ``take``. Dropped, every worker goes back to pulling the entire expired set + in one unbounded query at the same calendar boundary.""" + find_many = _run_reset_query(table_name, limit=7) + + assert find_many.await_args.kwargs["take"] == 7 + + +@pytest.mark.parametrize("table_name", ["user", "team", "budget", "key"]) +def test_get_data_reset_query_skips_rows_with_no_budget_duration(table_name): + """A row with a past budget_reset_at but no budget_duration has no next + window to move to, so it stays due forever. Fetching it means re-reading and + re-zeroing it on every tick, and a full chunk of such rows makes the paged + scan report no progress and starve the whole phase. + """ + find_many = _run_reset_query(table_name) + + assert find_many.await_args.kwargs["where"]["NOT"] == {"budget_duration": None} + + +@pytest.mark.parametrize("table_name", ["user", "team", "budget", "key"]) +def test_get_data_reset_query_is_unlimited_when_no_limit_is_passed(table_name): + """Callers that pass no limit keep the old unbounded behaviour.""" + find_many = _run_reset_query(table_name) + + assert find_many.await_args.kwargs.get("take") is None @pytest.mark.parametrize("table_name", ["user", "team"]) @@ -1838,8 +1608,327 @@ def test_get_data_reset_query_selects_null_budget_reset_at(table_name): setattr(getattr(client.db, table_attr), "find_many", find_many) now = datetime.now(timezone.utc) - asyncio.run( - client.get_data(table_name=table_name, query_type="find_all", reset_at=now) - ) + asyncio.run(client.get_data(table_name=table_name, query_type="find_all", reset_at=now)) _asserts_null_reset_is_due(_extract_reset_where(find_many)) + + +def _key_row(token: str, budget_duration: Any = "30d"): + """A key that is already due for a reset, shaped like a get_data() row.""" + now = datetime.now(timezone.utc) + return type( + "LiteLLM_VerificationToken", + (), + { + "spend": 100.0, + "budget_duration": budget_duration, + "budget_reset_at": now - timedelta(hours=1), + "token": token, + }, + ) + + +def _user_row(user_id: str, budget_duration: Any = "30d"): + now = datetime.now(timezone.utc) + return type( + "LiteLLM_UserTable", + (), + { + "spend": 100.0, + "budget_duration": budget_duration, + "budget_reset_at": now - timedelta(hours=1), + "user_id": user_id, + }, + ) + + +def _team_row(team_id: str, budget_duration: Any = "30d"): + now = datetime.now(timezone.utc) + return type( + "LiteLLM_TeamTable", + (), + { + "spend": 100.0, + "budget_duration": budget_duration, + "budget_reset_at": now - timedelta(hours=1), + "team_id": team_id, + }, + ) + + +# --------------------------------------------------------------------------- +# Chunked batches +# --------------------------------------------------------------------------- + + +class ChunkedPrismaClient(MockPrismaClient): + """Replays a scripted sequence of get_data chunks per table. + + The last chunk repeats forever, so a phase that fails to terminate keeps + seeing rows rather than quietly running out of data. + """ + + def __init__(self, chunks_by_table: Dict[str, List[List[Any]]]): + super().__init__() + self._chunks_by_table = chunks_by_table + self.fetches_by_table: Dict[str, int] = {} + + async def get_data(self, table_name, query_type, **kwargs): + self.get_data_calls.append({"table_name": table_name, "query_type": query_type, **kwargs}) + chunks = self._chunks_by_table.get(table_name) + if not chunks: + return [] + index = self.fetches_by_table.get(table_name, 0) + self.fetches_by_table[table_name] = index + 1 + return chunks[min(index, len(chunks) - 1)] + + +def _chunked_job(chunks_by_table): + client = ChunkedPrismaClient(chunks_by_table) + return client, ResetBudgetJob(proxy_logging_obj=MockProxyLogging(), prisma_client=client) + + +def _fetch_limits(client, table_name): + return [call.get("limit") for call in client.get_data_calls if call["table_name"] == table_name] + + +def test_key_reset_walks_the_due_rows_one_chunk_at_a_time(monkeypatch): + """Each chunk is fetched under a LIMIT and committed on its own batch, so a + large backlog never becomes one giant transaction.""" + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + client, job = _chunked_job({"key": [[_key_row("k1"), _key_row("k2")], [_key_row("k3")]]}) + + asyncio.run(job.reset_budget_for_litellm_keys()) + + assert client.fetches_by_table["key"] == 2 + assert _fetch_limits(client, "key") == [2, 2] + assert len(client.db.batchers) == 2 + assert all(batcher.committed for batcher in client.db.batchers) + assert [len(batcher.calls) for batcher in client.db.batchers] == [2, 1] + assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2", "k3"] + + +def test_key_reset_stops_after_a_chunk_shorter_than_the_batch_size(monkeypatch): + """Fewer rows than the limit means the backlog is drained, so no follow-up + query is worth issuing.""" + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 5) + client, job = _chunked_job({"key": [[_key_row("k1")]]}) + + asyncio.run(job.reset_budget_for_litellm_keys()) + + assert client.fetches_by_table["key"] == 1 + + +def test_key_reset_stops_when_a_full_chunk_advances_nothing(monkeypatch): + """A key with no budget_duration keeps its past budget_reset_at, so the very + same rows come back on the next fetch. Treating those writes as progress + would re-read that chunk until the iteration cap, every tick.""" + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + stuck_chunk = [_key_row("k1", budget_duration=None), _key_row("k2", budget_duration=None)] + client, job = _chunked_job({"key": [stuck_chunk]}) + + asyncio.run(job.reset_budget_for_litellm_keys()) + + assert client.fetches_by_table["key"] == 1 + + +def test_key_reset_stops_when_the_fetch_fails(monkeypatch): + """A phase whose query raises has made no progress; retrying it in a tight + loop would just hammer a struggling database.""" + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + client, job = _chunked_job({"key": [[_key_row("k1"), _key_row("k2")]]}) + + async def _boom(table_name, query_type, **kwargs): + client.get_data_calls.append({"table_name": table_name, "query_type": query_type, **kwargs}) + raise RuntimeError("db is down") + + client.get_data = _boom + + asyncio.run(job.reset_budget_for_litellm_keys()) + + assert len(client.get_data_calls) == 1 + + +def test_key_reset_is_capped_at_max_chunks_per_run(monkeypatch): + """Backstop against a phase that keeps making progress forever: the run ends + and the leftovers wait for the next tick.""" + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 1) + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", 3) + client, job = _chunked_job({"key": [[_key_row("k1")]]}) + + asyncio.run(job.reset_budget_for_litellm_keys()) + + assert client.fetches_by_table["key"] == 3 + + +@pytest.mark.parametrize( + "phase, table_name, row_factory", + [ + ("reset_budget_for_litellm_users", "user", lambda uid: _user_row(uid)), + ("reset_budget_for_litellm_teams", "team", lambda tid: _team_row(tid)), + ], + ids=["users", "teams"], +) +def test_user_and_team_resets_are_chunked_too(monkeypatch, phase, table_name, row_factory): + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + client, job = _chunked_job({table_name: [[row_factory("a"), row_factory("b")], [row_factory("c")]]}) + + asyncio.run(getattr(job, phase)()) + + assert client.fetches_by_table[table_name] == 2 + assert _fetch_limits(client, table_name) == [2, 2] + assert len(client.db.batchers) == 2 + assert len(_batch_writes(client, table_name, op="update")) == 3 + + +def test_budget_table_reset_walks_chunks_until_it_runs_dry(monkeypatch): + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + client, job = _chunked_job({"budget": [[_budget_row("b1"), _budget_row("b2")], [_budget_row("b3")]]}) + + asyncio.run(job.reset_budget_for_litellm_budget_table()) + + assert client.fetches_by_table["budget"] == 2 + assert _fetch_limits(client, "budget") == [2, 2] + assert len(client.db.batchers) == 2 + assert all(batcher.committed for batcher in client.db.batchers) + assert [w["where"]["budget_id"] for w in _batch_writes(client, "budget", op="update_many")] == ["b1", "b2", "b3"] + + +def test_budget_table_reset_stops_when_a_full_chunk_advances_no_window(monkeypatch): + """A tier with no budget_duration has its linked spend zeroed but keeps its + past budget_reset_at, so it stays due. Counting those spend writes as + progress would re-read the same chunk until the cap.""" + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + stuck_chunk = [_budget_row("b1", budget_duration=None), _budget_row("b2", budget_duration=None)] + client, job = _chunked_job({"budget": [stuck_chunk]}) + + asyncio.run(job.reset_budget_for_litellm_budget_table()) + + assert client.fetches_by_table["budget"] == 1 + assert _batch_writes(client, "budget", op="update_many") == [] + + +def test_budget_table_reset_stops_when_the_cascade_fails(monkeypatch): + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + client, job = _chunked_job({"budget": [[_budget_row("b1"), _budget_row("b2")]]}) + client.db = FailingCommitDB() + + asyncio.run(job.reset_budget_for_litellm_budget_table()) + + assert client.fetches_by_table["budget"] == 1 + + +# --------------------------------------------------------------------------- +# Progress means "no longer due", not "was written" +# --------------------------------------------------------------------------- + + +def test_key_reset_stops_when_the_new_reset_time_is_not_in_the_future(monkeypatch): + """A "0s" budget_duration resolves to the current time, so the row is written + and comes straight back on the next fetch. Treating a written row as progress + burns the whole per-run chunk cap on rows that never move. + """ + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + stuck_chunk = [_key_row("k1", budget_duration="0s"), _key_row("k2", budget_duration="0s")] + client, job = _chunked_job({"key": [stuck_chunk]}) + + asyncio.run(job.reset_budget_for_litellm_keys()) + + assert client.fetches_by_table["key"] == 1 + assert len(_batch_writes(client, "key", op="update")) == 2 + + +def test_budget_table_reset_stops_when_the_new_window_is_not_in_the_future(monkeypatch): + """Same zero-length window on the budget tier: advancing it to now leaves it + due, so the cascade must not report progress.""" + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + stuck_chunk = [_budget_row("b1", budget_duration="0s"), _budget_row("b2", budget_duration="0s")] + client, job = _chunked_job({"budget": [stuck_chunk]}) + + asyncio.run(job.reset_budget_for_litellm_budget_table()) + + assert client.fetches_by_table["budget"] == 1 + assert len(_batch_writes(client, "budget", op="update_many")) == 2 + + +class PoisonRow: + """A row the in-memory reset cannot write, like the DataError rows in #27730.""" + + token = "poison" + budget_duration = "30d" + budget_reset_at = None + + def __setattr__(self, name: str, value: Any) -> None: + raise RuntimeError("simulated failure resetting this row") + + +class RecordingServiceLogging: + def __init__(self): + self.success_calls: List[Dict[str, Any]] = [] + self.failure_calls: List[Dict[str, Any]] = [] + + async def async_service_success_hook(self, **kwargs): + self.success_calls.append(kwargs) + + async def async_service_failure_hook(self, **kwargs): + self.failure_calls.append(kwargs) + + +class RecordingProxyLogging: + def __init__(self): + self.service_logging_obj = RecordingServiceLogging() + + +def _run_and_drain_hooks(make_coro): + """The service hooks are fired as tasks; give them a turn before asserting.""" + + async def _run(): + await make_coro() + await asyncio.sleep(0.05) + + asyncio.run(_run()) + + +def test_key_reset_keeps_paging_when_some_rows_in_a_chunk_fail(monkeypatch): + """One row that cannot be reset must not cost the phase its remaining chunks: + the rows that did reset are committed and are real progress, and the failure + is reported instead of aborting the run. + """ + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + client = ChunkedPrismaClient({"key": [[PoisonRow(), _key_row("k1")], [_key_row("k2")]]}) + logging_obj = RecordingProxyLogging() + job = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=client) + + _run_and_drain_hooks(job.reset_budget_for_litellm_keys) + + assert client.fetches_by_table["key"] == 2 + assert [w["where"]["token"] for w in _batch_writes(client, "key", op="update")] == ["k1", "k2"] + assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == ["reset_budget_keys"] + assert set(logging_obj.service_logging_obj.failure_calls[0]["event_metadata"]) == { + "num_keys_found", + "keys_found", + } + assert [call["call_type"] for call in logging_obj.service_logging_obj.success_calls] == ["reset_budget_keys"] + + +@pytest.mark.parametrize( + "phase, table_name, row_factory, call_type", + [ + ("reset_budget_for_litellm_users", "user", _user_row, "reset_budget_users"), + ("reset_budget_for_litellm_teams", "team", _team_row, "reset_budget_teams"), + ], + ids=["users", "teams"], +) +def test_user_and_team_chunks_report_progress_despite_a_failed_row( + monkeypatch, phase, table_name, row_factory, call_type +): + monkeypatch.setattr(reset_budget_job_module, "RESET_BUDGET_JOB_BATCH_SIZE", 2) + client = ChunkedPrismaClient({table_name: [[PoisonRow(), row_factory("a")], [row_factory("b")]]}) + logging_obj = RecordingProxyLogging() + job = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=client) + + _run_and_drain_hooks(getattr(job, phase)) + + assert client.fetches_by_table[table_name] == 2 + assert len(_batch_writes(client, table_name, op="update")) == 2 + assert [call["call_type"] for call in logging_obj.service_logging_obj.failure_calls] == [call_type] diff --git a/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py new file mode 100644 index 00000000000..c2d0f64461a --- /dev/null +++ b/tests/test_litellm/proxy/db/test_daily_spend_bulk_upsert.py @@ -0,0 +1,187 @@ +"""Tests for the single-statement daily spend upsert (LIT-5291).""" + +import re + +import pytest + +from litellm.proxy.db.daily_spend_bulk_upsert import ( + DAILY_SPEND_TABLES, + build_bulk_upsert, + conflict_key, + merge_by_conflict_key, +) +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + +TAG_TABLE = DAILY_SPEND_TABLES["tag"] +USER_TABLE = DAILY_SPEND_TABLES["user"] + +# Every nullable member of the unique constraint, so a test that only varied the provider +# cannot pass while a sibling column still leaks a NULL into the conflict target. +NULLABLE_KEY_COLUMNS = ("model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint") + + +def tag_txn(**overrides): + return { + "tag": "team-a", + "date": "2026-08-10", + "api_key": "sk-hash", + "model": "gpt-4o-mini", + "model_group": "gpt-4o-mini", + "custom_llm_provider": "openai", + "mcp_namespaced_tool_name": "", + "endpoint": "/chat/completions", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.25, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + "request_id": "req-1", + **overrides, + } + + +@pytest.mark.parametrize("column", NULLABLE_KEY_COLUMNS) +def test_conflict_key_normalizes_every_nullable_key_column(column): + """A NULL member can never match itself in a unique index, so the row would be + re-inserted on every flush. Each nullable key column must arrive as ''.""" + key = conflict_key(TAG_TABLE, tag_txn(**{column: None})) + + assert "" in key + assert None not in key + assert key == conflict_key(TAG_TABLE, tag_txn(**{column: ""})) + + +@pytest.mark.parametrize("order", [("null_first"), ("empty_first")]) +def test_null_and_empty_provider_merge_into_one_row(order): + """Two queue entries differing only in NULL versus '' arbitrate to the same row. + Postgres rejects one statement touching a row twice, so they must be folded first. + Asserted under both input orders: a single ordering would prove nothing here.""" + null_entry = tag_txn(custom_llm_provider=None, spend=0.25, api_requests=1) + empty_entry = tag_txn(custom_llm_provider="", spend=0.75, api_requests=3) + transactions = (null_entry, empty_entry) if order == "null_first" else (empty_entry, null_entry) + + merged = merge_by_conflict_key(TAG_TABLE, transactions) + + assert len(merged) == 1 + _, folded = merged[0] + assert folded["spend"] == pytest.approx(1.0) + assert folded["api_requests"] == 4 + + +def test_distinct_keys_are_not_merged_and_are_ordered_deterministically(): + unordered = (tag_txn(tag="z-team"), tag_txn(tag="a-team"), tag_txn(tag="m-team")) + + merged = merge_by_conflict_key(TAG_TABLE, unordered) + + assert [txn["tag"] for _, txn in merged] == ["a-team", "m-team", "z-team"] + assert merged == merge_by_conflict_key(TAG_TABLE, tuple(reversed(unordered))) + + +def test_one_statement_carries_every_row_in_the_batch(): + batch = merge_by_conflict_key(TAG_TABLE, tuple(tag_txn(tag=f"team-{i}") for i in range(100))) + + sql, params = build_bulk_upsert(TAG_TABLE, batch) + + assert sql.count("INSERT INTO") == 1 + assert len(re.findall(r"ON CONFLICT", sql)) == 1 + # 22 bound columns per row plus the inlined updated_at, so the row count is what + # separates one multi-row statement from a hundred single-row ones. + assert len(params) == 100 * 22 + assert "$2200::text" in sql + assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1 + + +def test_conflict_target_is_the_full_unique_constraint(): + sql, _ = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(),))) + + conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql) + assert conflict_target is not None + assert conflict_target.group(1) == ( + '"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint"' + ) + + +@pytest.mark.parametrize( + "column", + ["prompt_tokens", "completion_tokens", "spend", "api_requests", "successful_requests", "failed_requests"], +) +def test_counters_increment_rather_than_overwrite(column): + """An overwrite would silently discard every earlier flush's spend for that row.""" + sql, _ = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(),))) + + assert f'"{column}" = "LiteLLM_DailyTagSpend"."{column}" + EXCLUDED."{column}"' in sql + + +def test_request_id_is_preserved_when_a_later_batch_carries_none(): + sql, params = build_bulk_upsert( + TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(request_id=None),)) + ) + + assert '"request_id" = COALESCE(EXCLUDED."request_id", "LiteLLM_DailyTagSpend"."request_id")' in sql + assert None in params + + +def test_non_tag_tables_carry_no_request_id_column(): + user_txn = {**tag_txn(), "user_id": "u-1"} + del user_txn["tag"] + + sql, _ = build_bulk_upsert(USER_TABLE, merge_by_conflict_key(USER_TABLE, (user_txn,))) + + assert "request_id" not in sql + assert '"user_id"' in sql + + +class _RecordingDb: + def __init__(self) -> None: + self.statements: list[tuple[str, tuple[object, ...]]] = [] + + async def execute_raw(self, query: str, *args: object) -> int: + self.statements.append((query, args)) + return len(args) + + +class _RecordingPrismaClient: + def __init__(self) -> None: + self.db = _RecordingDb() + + +@pytest.mark.asyncio +async def test_writer_issues_one_statement_per_batch_not_one_per_key(): + """The whole point of LIT-5291: 250 aggregated keys must not become 250 statements.""" + prisma_client = _RecordingPrismaClient() + transactions = {f"k{i}": tag_txn(tag=f"team-{i}") for i in range(250)} + + await DBSpendUpdateWriter.update_daily_tag_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=None, + daily_spend_transactions=transactions, + ) + + # 250 keys at a batch size of 100 is three statements, one per batch. + assert len(prisma_client.db.statements) == 3 + assert [statement.count("ON CONFLICT") for statement, _ in prisma_client.db.statements] == [1, 1, 1] + assert transactions == {} + + +@pytest.mark.asyncio +async def test_writer_survives_a_transaction_whose_key_columns_are_null(): + """A NULL key column used to raise out of prisma and drop the whole batch's spend.""" + prisma_client = _RecordingPrismaClient() + transactions = { + "mcp": tag_txn(model=None, custom_llm_provider=None, mcp_namespaced_tool_name="server/tool"), + "chat": tag_txn(), + } + + await DBSpendUpdateWriter.update_daily_tag_spend( + n_retry_times=0, + prisma_client=prisma_client, + proxy_logging_obj=None, + daily_spend_transactions=transactions, + ) + + assert len(prisma_client.db.statements) == 1 + _, params = prisma_client.db.statements[0] + assert None not in params[:9] + assert transactions == {} diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 191080e3a48..51810a28cdd 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2,6 +2,7 @@ import asyncio import copy import json import os +import re import sys sys.path.insert( @@ -9,6 +10,7 @@ sys.path.insert( ) # Adds the parent directory to the system path +from collections.abc import Callable from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, call, patch @@ -232,21 +234,49 @@ async def test_update_database_skips_tool_usage_when_spend_logs_disabled(): assert prisma.tool_usage_transactions == [] +Statement = tuple[str, tuple[object, ...]] + + +class _RecordingDb: + """Records the statements the writer sends, in place of a real query engine.""" + + def __init__(self, execute_raw: Callable[[], int] | None = None) -> None: + self.statements: list[Statement] = [] + self._execute_raw = execute_raw + + async def execute_raw(self, query: str, *args: object) -> int: + self.statements.append((query, args)) + if self._execute_raw is not None: + return self._execute_raw() + return len(args) + + +class _RecordingPrisma: + def __init__(self, execute_raw: Callable[[], int] | None = None) -> None: + self.db = _RecordingDb(execute_raw=execute_raw) + + +def _row_values(statement: Statement, column: str) -> list[object]: + """Every row's value for one column, read out of the flat parameter tuple.""" + sql, params = statement + header = re.search(r"INSERT INTO \"[A-Za-z_]+\" \(([^)]*)\)", sql) + assert header is not None, sql + columns = header.group(1).split(", ") + stride = len(columns) - 1 # updated_at is inlined, not bound + offset = columns.index(f'"{column}"') + return [params[row * stride + offset] for row in range(len(params) // stride)] + + @pytest.mark.asyncio async def test_update_daily_spend_with_null_entity_id(): """ - Test that table.upsert is called even when entity_id is null + A null entity_id must still be written, so the 'global view' keeps that spend. - Ensures 'global view' has all daily spend transactions + It is stored as '' rather than NULL: a NULL can never match itself in the unique + index, so such a row would be re-inserted on every flush instead of aggregating. """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailyuserspend = mock_table + prisma_client = _RecordingPrisma() - # Create a transaction with null entity_id daily_spend_transactions = { "test_key": { "user_id": None, # null entity_id @@ -263,49 +293,30 @@ async def test_update_daily_spend_with_null_entity_id(): } } - # Call the method await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - # Verify that table.upsert was called - mock_table.upsert.assert_called_once() - - # Verify the where clause contains null entity_id - call_args = mock_table.upsert.call_args[1] - where_clause = call_args["where"][ - "user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" - ] - assert where_clause["user_id"] is None - assert where_clause["date"] == "2024-01-01" - assert where_clause["api_key"] == "test-api-key" - assert where_clause["model"] == "gpt-4" - assert where_clause["custom_llm_provider"] == "openai" - assert where_clause["mcp_namespaced_tool_name"] == "" - assert where_clause["endpoint"] == "" - - # Verify the create data contains null entity_id - create_data = call_args["data"]["create"] - assert create_data["user_id"] is None - assert create_data["date"] == "2024-01-01" - assert create_data["api_key"] == "test-api-key" - assert create_data["model"] == "gpt-4" - assert create_data["custom_llm_provider"] == "openai" - assert create_data["mcp_namespaced_tool_name"] == "" - assert create_data["endpoint"] == "" - assert create_data["prompt_tokens"] == 10 - assert create_data["completion_tokens"] == 20 - assert create_data["spend"] == 0.1 - assert create_data["api_requests"] == 1 - assert create_data["successful_requests"] == 1 - assert create_data["failed_requests"] == 0 + assert len(prisma_client.db.statements) == 1 + statement = prisma_client.db.statements[0] + assert _row_values(statement, "user_id") == [""] + assert _row_values(statement, "date") == ["2024-01-01"] + assert _row_values(statement, "api_key") == ["test-api-key"] + assert _row_values(statement, "model") == ["gpt-4"] + assert _row_values(statement, "custom_llm_provider") == ["openai"] + assert _row_values(statement, "mcp_namespaced_tool_name") == [""] + assert _row_values(statement, "endpoint") == [""] + assert _row_values(statement, "prompt_tokens") == [10] + assert _row_values(statement, "completion_tokens") == [20] + assert _row_values(statement, "spend") == [0.1] + assert _row_values(statement, "api_requests") == [1] + assert _row_values(statement, "successful_requests") == [1] + assert _row_values(statement, "failed_requests") == [0] def _daily_txn(user_id: str = "user1") -> dict: @@ -333,24 +344,24 @@ async def test_update_daily_spend_does_not_retry_post_send_ambiguous_errors(): # batch (loudly), never retry it. import httpx - mock_prisma_client = MagicMock() - mock_prisma_client.db.batch_ = MagicMock(side_effect=httpx.ReadTimeout("ambiguous")) + def raise_read_timeout(): + raise httpx.ReadTimeout("ambiguous") + + prisma_client = _RecordingPrisma(execute_raw=raise_read_timeout) proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() with pytest.raises(httpx.ReadTimeout): await DBSpendUpdateWriter._update_daily_spend( n_retry_times=3, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=proxy_logging, daily_spend_transactions={"k1": _daily_txn()}, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - mock_prisma_client.db.batch_.assert_called_once() + assert len(prisma_client.db.statements) == 1 @pytest.mark.asyncio @@ -359,12 +370,15 @@ async def test_update_daily_spend_retries_connect_errors(monkeypatch): # the one failure the writer may retry. import httpx - mock_batcher = MagicMock() - good_ctx = MagicMock() - good_ctx.__aenter__ = AsyncMock(return_value=mock_batcher) - good_ctx.__aexit__ = AsyncMock(return_value=None) - mock_prisma_client = MagicMock() - mock_prisma_client.db.batch_ = MagicMock(side_effect=[httpx.ConnectError("down"), good_ctx]) + outcomes = iter([httpx.ConnectError("down"), None]) + + def first_attempt_disconnects(): + outcome = next(outcomes) + if outcome is not None: + raise outcome + return 1 + + prisma_client = _RecordingPrisma(execute_raw=first_attempt_disconnects) proxy_logging = MagicMock() proxy_logging.failure_handler = AsyncMock() @@ -374,16 +388,14 @@ async def test_update_daily_spend_retries_connect_errors(monkeypatch): monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", fake_sleep) await DBSpendUpdateWriter._update_daily_spend( n_retry_times=3, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=proxy_logging, daily_spend_transactions={"k1": _daily_txn()}, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - assert mock_prisma_client.db.batch_.call_count == 2 + assert len(prisma_client.db.statements) == 2 @pytest.mark.asyncio @@ -393,19 +405,12 @@ async def test_update_daily_spend_sorting(): Ensures that writes are sorted between transactions to minimize deadlocks """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailyuserspend = mock_table + prisma_client = _RecordingPrisma() - # Create a 50 transactions with out-of-order entity_ids - # In reality we sort using multiple fields, but entity_id is sufficient to test sorting - daily_spend_transactions = {} - upsert_calls = [] - for i in range(50): - daily_spend_transactions[f"test_key_{i}"] = { + # 50 transactions with out-of-order entity_ids. In reality we sort using multiple + # fields, but entity_id is sufficient to test sorting. + daily_spend_transactions = { + f"test_key_{i}": { "user_id": f"user{60-i}", # user60 ... user11, reverse order "date": "2024-01-01", "api_key": "test-api-key", @@ -418,63 +423,22 @@ async def test_update_daily_spend_sorting(): "successful_requests": 1, "failed_requests": 0, } - upsert_calls.append( - call( - where={ - "user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": { - "user_id": f"user{i+11}", # user11 ... user60, sorted order - "date": "2024-01-01", - "api_key": "test-api-key", - "model": "gpt-4", - "custom_llm_provider": "openai", - "mcp_namespaced_tool_name": "", - "endpoint": "", - } - }, - data={ - "create": { - "user_id": f"user{i+11}", - "date": "2024-01-01", - "api_key": "test-api-key", - "model": "gpt-4", - "model_group": None, - "mcp_namespaced_tool_name": "", - "custom_llm_provider": "openai", - "endpoint": "", - "prompt_tokens": 10, - "completion_tokens": 20, - "spend": 0.1, - "api_requests": 1, - "successful_requests": 1, - "failed_requests": 0, - }, - "update": { - "prompt_tokens": {"increment": 10}, - "completion_tokens": {"increment": 20}, - "spend": {"increment": 0.1}, - "api_requests": {"increment": 1}, - "successful_requests": {"increment": 1}, - "failed_requests": {"increment": 0}, - "endpoint": "", - }, - }, - ) - ) + for i in range(50) + } - # Call the method await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - # Verify that table.upsert was called - mock_table.upsert.assert_has_calls(upsert_calls) + assert len(prisma_client.db.statements) == 1 + written = _row_values(prisma_client.db.statements[0], "user_id") + assert written == sorted(written) + assert written[0] == "user11" and written[-1] == "user60" @pytest.mark.asyncio @@ -485,11 +449,7 @@ async def test_update_daily_spend_drains_all_batches_over_batch_size(): only the first 100 sorted items were upserted then the method returned, silently dropping the remaining entities. """ - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailyuserspend = mock_table + prisma_client = _RecordingPrisma() num_entities = 250 daily_spend_transactions = { @@ -511,17 +471,16 @@ async def test_update_daily_spend_drains_all_batches_over_batch_size(): await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - assert mock_table.upsert.call_count == num_entities - assert mock_prisma_client.db.batch_.call_count == 3 + assert len(prisma_client.db.statements) == 3 + all_written = [uid for statement in prisma_client.db.statements for uid in _row_values(statement, "user_id")] + assert sorted(all_written) == sorted(f"user{i:04d}" for i in range(num_entities)) assert daily_spend_transactions == {} @@ -530,14 +489,8 @@ async def test_update_daily_spend_tag_with_request_id(): """ Test that request_id is included in update_data when updating tag transactions. """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailytagspend = mock_table + prisma_client = _RecordingPrisma() - # Create a transaction with request_id daily_spend_transactions = { "test_key": { "tag": "prod-tag", @@ -556,26 +509,19 @@ async def test_update_daily_spend_tag_with_request_id(): } } - # Call the method await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="tag", entity_id_field="tag", - table_name="litellm_dailytagspend", - unique_constraint_name="tag_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name", ) - # Verify that table.upsert was called - mock_table.upsert.assert_called_once() - - # Verify request_id is in update_data - call_args = mock_table.upsert.call_args[1] - update_data = call_args["data"]["update"] - assert "request_id" in update_data - assert update_data["request_id"] == "test-request-id-123" + assert len(prisma_client.db.statements) == 1 + sql, _ = prisma_client.db.statements[0] + assert _row_values(prisma_client.db.statements[0], "request_id") == ["test-request-id-123"] + assert '"request_id" = COALESCE(EXCLUDED."request_id", "LiteLLM_DailyTagSpend"."request_id")' in sql @pytest.mark.asyncio @@ -587,12 +533,7 @@ async def test_update_daily_spend_with_none_values_in_sorting_fields(): are None, the sorting doesn't crash with TypeError: '<' not supported between instances of 'NoneType' and 'str'. """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher - mock_batcher.litellm_dailyuserspend = mock_table + prisma_client = _RecordingPrisma() # Create transactions with None values in various sorting fields daily_spend_transactions = { @@ -666,17 +607,20 @@ async def test_update_daily_spend_with_none_values_in_sorting_fields(): # Call the method - this should not raise TypeError await DBSpendUpdateWriter._update_daily_spend( n_retry_times=1, - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=MagicMock(), daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) - # Verify that table.upsert was called (should be called 5 times, once for each transaction) - assert mock_table.upsert.call_count == 5 + # All five distinct rows are written, in one statement, with no NULL anywhere in + # the conflict key. + assert len(prisma_client.db.statements) == 1 + statement = prisma_client.db.statements[0] + assert len(_row_values(statement, "user_id")) == 5 + for column in ("user_id", "date", "api_key", "model", "custom_llm_provider"): + assert None not in _row_values(statement, column) # Tag Spend Tracking Tests @@ -1384,19 +1328,10 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): """ from litellm._logging import verbose_proxy_logger - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_batch_context = MagicMock() - mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher) - mock_batcher.litellm_dailyuserspend = mock_table + def raise_constraint_violation(): + raise Exception("Unique constraint violation") - # Make the batch context manager's exit raise an exception - # This simulates a batch commit failure (e.g., unique constraint violation) - test_exception = Exception("Unique constraint violation") - mock_batch_context.__aexit__ = AsyncMock(side_effect=test_exception) - mock_prisma_client.db.batch_.return_value = mock_batch_context + prisma_client = _RecordingPrisma(execute_raw=raise_constraint_violation) # Create a transaction daily_spend_transactions = { @@ -1427,28 +1362,22 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure(): with pytest.raises(Exception, match="Unique constraint violation"): await DBSpendUpdateWriter._update_daily_spend( n_retry_times=0, # No retries to make test faster - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=mock_proxy_logging, daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) # Verify that the error was logged with detailed information. # spend_log_error formats the message via ``%`` interpolation, so # render the call args before asserting on substrings. assert mock_error_logger.called - call = mock_error_logger.call_args - formatted = call.args[0] % call.args[1:] + logged = mock_error_logger.call_args + formatted = logged.args[0] % logged.args[1:] assert "Daily user spend batch upsert failed" in formatted - assert "Table: litellm_dailyuserspend" in formatted - assert ( - "Constraint: user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" - in formatted - ) - assert "Batch size: 1" in formatted + assert "Table: LiteLLM_DailyUserSpend" in formatted + assert "Rows: 1" in formatted assert "Unique constraint violation" in formatted @@ -1458,13 +1387,10 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): Test that when batch upsert fails, the exception is properly re-raised after logging. This ensures that error handling continues to work correctly upstream. """ - # Setup - mock_prisma_client = MagicMock() - mock_batcher = MagicMock() - mock_table = MagicMock() - mock_batch_context = MagicMock() - mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher) - mock_batcher.litellm_dailyuserspend = mock_table + def raise_connection_lost(): + raise ValueError("Database connection lost") + + prisma_client = _RecordingPrisma(execute_raw=raise_connection_lost) # Create a transaction daily_spend_transactions = { @@ -1483,11 +1409,6 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): } } - # Create a custom exception to verify it's re-raised - custom_exception = ValueError("Database connection lost") - mock_batch_context.__aexit__ = AsyncMock(side_effect=custom_exception) - mock_prisma_client.db.batch_.return_value = mock_batch_context - # Create a mock proxy_logging_obj with failure_handler as AsyncMock mock_proxy_logging = MagicMock() mock_proxy_logging.failure_handler = AsyncMock() @@ -1496,13 +1417,11 @@ async def test_update_daily_spend_re_raises_exception_after_logging(): with pytest.raises(ValueError, match="Database connection lost"): await DBSpendUpdateWriter._update_daily_spend( n_retry_times=0, # No retries to make test faster - prisma_client=mock_prisma_client, + prisma_client=prisma_client, proxy_logging_obj=mock_proxy_logging, daily_spend_transactions=daily_spend_transactions, entity_type="user", entity_id_field="user_id", - table_name="litellm_dailyuserspend", - unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", ) @@ -2196,9 +2115,11 @@ async def test_daily_transaction_carries_compression_saved_tokens(): model_info = litellm.get_model_info(model="claude-sonnet-5", custom_llm_provider="anthropic") input_cost = model_info["input_cost_per_token"] or 0.0 cache_read_cost = model_info.get("cache_read_input_token_cost") or input_cost + cache_write_cost = model_info.get("cache_creation_input_token_cost") or input_cost assert transaction["compression_savings_spend"] == pytest.approx(7600 * input_cost) assert transaction["prompt_caching_savings_spend"] == pytest.approx( 40 * max(input_cost - cache_read_cost, 0.0) + - 15 * (cache_write_cost - input_cost) ) assert transaction["compression_savings_spend"] > 0 assert transaction["prompt_caching_savings_spend"] > 0 diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 1c29c287c3a..076151fcd3b 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3,6 +3,7 @@ Unit Tests for the max parallel request limiter v3 for the proxy """ import asyncio +import logging import os import sys import time @@ -5100,3 +5101,456 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): f"reservation pass, got: {response}" ) assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"] + + +STATIC_OUTPUT_FLOOR = 1024 +ONE_TOKEN_PROMPT = [{"role": "user", "content": "hello"}] +ONE_TOKEN_PROMPT_INPUT_ESTIMATE = 1 + + +async def _reserved_tokens_for( + handler, + local_cache, + user_api_key_dict, + data, + call_type="completion", +): + """Drive the pre-call hook and read back what landed on the :tokens counter.""" + tokens_key = handler.create_rate_limit_keys( + key="api_key", value=user_api_key_dict.api_key, rate_limit_type="tokens" + ) + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + return int(await local_cache.async_get_cache(key=tokens_key) or 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_metadata, team_metadata, expected_output_estimate, tier", + [ + ( + { + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 3001}, + "default_estimated_output_tokens": 2002, + }, + { + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 1503}, + "default_estimated_output_tokens": 777, + }, + 3001, + "key per-model wins over every other tier", + ), + ( + {"default_estimated_output_tokens": 2002}, + { + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 1503}, + "default_estimated_output_tokens": 777, + }, + 2002, + "key global wins over team config", + ), + ( + {"default_estimated_output_tokens_per_model": {"some-other-model": 9999}}, + { + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 1503}, + "default_estimated_output_tokens": 777, + }, + 1503, + "team per-model wins when the key has no applicable entry", + ), + ( + {}, + {"default_estimated_output_tokens": 777}, + 777, + "team global is the last configured tier", + ), + ({}, {}, STATIC_OUTPUT_FLOOR, "unconfigured falls back to the static floor"), + ( + {"unrelated": "value"}, + {"unrelated": "value"}, + STATIC_OUTPUT_FLOOR, + "unrelated metadata changes nothing", + ), + ( + {"default_estimated_output_tokens": "not-a-number"}, + {}, + STATIC_OUTPUT_FLOOR, + "malformed config falls back to the static floor instead of erroring", + ), + ( + {"default_estimated_output_tokens": 0}, + {}, + STATIC_OUTPUT_FLOOR, + "a non-positive estimate is rejected, not reserved", + ), + ], +) +async def test_estimated_output_tokens_resolution_precedence( + monkeypatch, key_metadata, team_metadata, expected_output_estimate, tier +): + """The no-max_tokens output reservation resolves per key / team / model. + + Every configured value here is distinct from the static 1024 floor and + from the input estimate, so the reserved amount identifies which tier the + resolver picked. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token(f"sk-estimate-{expected_output_estimate}-{tier}"), + tpm_limit=1_000_000, + metadata=key_metadata, + team_metadata=team_metadata, + ) + + reserved = await _reserved_tokens_for( + handler, + local_cache, + user_api_key_dict, + {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + ) + + assert reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + expected_output_estimate, tier + + +@pytest.mark.asyncio +async def test_request_max_tokens_outranks_configured_estimate(monkeypatch): + """An explicit request-level max_tokens stays the top of the precedence order.""" + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-estimate-explicit-max-tokens"), + tpm_limit=1_000_000, + metadata={"default_estimated_output_tokens": 2002}, + ) + + reserved = await _reserved_tokens_for( + handler, + local_cache, + user_api_key_dict, + {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT, "max_tokens": 42}, + ) + + assert reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 42 + + +@pytest.mark.asyncio +async def test_configured_estimate_does_not_apply_to_embeddings(monkeypatch): + """Embeddings generate no output, so a declared output estimate must not be reserved.""" + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-estimate-embeddings"), + tpm_limit=1_000_000, + metadata={"default_estimated_output_tokens": 2002}, + ) + + reserved = await _reserved_tokens_for( + handler, + local_cache, + user_api_key_dict, + {"model": "text-embedding-3-small", "input": "hello"}, + call_type="embeddings", + ) + + assert reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + + +@pytest.mark.asyncio +async def test_configured_estimate_applies_to_contentless_requests(monkeypatch): + """A declared estimate describes generation, so it holds even with no prompt body. + + Without config such a request reserves the 1-token floor only; the + declaration is what makes concurrent tool-call continuations countable. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + configured = UserAPIKeyAuth( + api_key=hash_token("sk-estimate-contentless-configured"), + tpm_limit=1_000_000, + metadata={"default_estimated_output_tokens": 2002}, + ) + unconfigured = UserAPIKeyAuth( + api_key=hash_token("sk-estimate-contentless-plain"), + tpm_limit=1_000_000, + ) + + assert ( + await _reserved_tokens_for( + handler, local_cache, configured, {"model": "gpt-4o-mini", "messages": []} + ) + == 2002 + ) + assert ( + await _reserved_tokens_for( + handler, local_cache, unconfigured, {"model": "gpt-4o-mini", "messages": []} + ) + == 1 + ) + + +@pytest.mark.asyncio +async def test_declared_estimate_never_tightens_the_small_tpm_clamp(monkeypatch): + """The small-TPM clamp can only be loosened by a declaration, never tightened. + + That clamp is the one place the proxy rewrites the caller's generation + budget, and it only fires below a 4096 TPM limit. A declaration above it + raises it, so the tenant is not truncated below what they said their + model emits; a declaration below it changes nothing, because an estimate + describes the typical response and must not become a hard cap that + truncates the tail. The reservation tracks whatever the clamp settles on, + so a small tenant can never generate more than was reserved. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + raised_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + raised_reserved = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-hard-cap-raised"), + tpm_limit=2000, + metadata={"default_estimated_output_tokens": 900}, + ), + raised_data, + ) + assert raised_data["max_tokens"] == 900 + assert raised_reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 900 + + lowered_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + lowered_reserved = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-hard-cap-lowered"), + tpm_limit=2000, + metadata={"default_estimated_output_tokens": 120}, + ), + lowered_data, + ) + assert lowered_data["max_tokens"] == 500 + assert lowered_reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 500 + + unconfigured_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + unconfigured_reserved = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-hard-cap-plain"), + tpm_limit=2000, + ), + unconfigured_data, + ) + assert unconfigured_data["max_tokens"] == 500 + assert unconfigured_reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 500 + + +@pytest.mark.asyncio +async def test_one_malformed_estimate_field_does_not_discard_the_other(monkeypatch): + """Each declared field is validated on its own. + + A per-model map with a bad entry must not take a valid global estimate + down with it, and a bad global must not hide a valid per-model entry. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + broken_map = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-broken-map"), + tpm_limit=1_000_000, + metadata={ + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": "huge"}, + "default_estimated_output_tokens": 2002, + }, + ), + {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + ) + assert broken_map == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 2002 + + broken_global = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-broken-global"), + tpm_limit=1_000_000, + metadata={ + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 3001}, + "default_estimated_output_tokens": -5, + }, + ), + {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + ) + assert broken_global == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 3001 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("declared", [100_000, 5000]) +async def test_declared_estimate_over_the_tpm_budget_is_honored_and_explained(monkeypatch, caplog, declared): + """A declaration bigger than the budget must not be silently shrunk. + + Capping it against the TPM limit would re-admit exactly the traffic this + feature exists to hold back, so the request is refused instead and the + reservation is explained rather than leaving an unexplained 429 loop. + + ``declared == tpm_limit`` is the boundary case: the declaration alone + equals the limit, so only adding the input estimate tips the reservation + over. Comparing the declaration against the limit rather than the + reservation would refuse this request while saying nothing. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token(f"sk-estimate-over-budget-{declared}"), + tpm_limit=5000, + metadata={"default_estimated_output_tokens": declared}, + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + explained = [ + record.getMessage() + for record in caplog.records + if "cannot be admitted even against an empty window" in record.getMessage() + ] + assert len(explained) == 1, f"expected exactly one explanation, got {explained}" + assert str(declared) in explained[0] + assert str(ONE_TOKEN_PROMPT_INPUT_ESTIMATE + declared) in explained[0] + assert "5000" in explained[0] + + +@pytest.mark.asyncio +async def test_a_key_that_declared_nothing_is_never_blamed_for_a_declaration(monkeypatch, caplog): + """A request can outgrow its budget on prompt size alone, with no declaration. + + The heuristic path reserves input plus the injected clamp, so a long + prompt against a small limit is refused without anyone having declared + anything. Blaming the declared field there would point an operator at a + setting they never set, to fix a 429 whose real cause is prompt size. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key=hash_token("sk-undeclared-long-prompt"), + tpm_limit=1000, + ), + cache=local_cache, + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "x" * 3600}]}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + assert not [ + record for record in caplog.records if "cannot be admitted even against an empty window" in record.getMessage() + ] + + +@pytest.mark.asyncio +async def test_declared_estimate_inside_the_tpm_budget_is_not_explained(monkeypatch, caplog): + """The explanation is for requests that cannot fit, not for every request. + + Without this, a correctly configured key would emit one line per call. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key=hash_token("sk-estimate-within-budget"), + tpm_limit=5000, + metadata={"default_estimated_output_tokens": 1000}, + ), + cache=local_cache, + data={"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + call_type="completion", + ) + + assert not [ + record for record in caplog.records if "cannot be admitted even against an empty window" in record.getMessage() + ] + + +@pytest.mark.asyncio +async def test_configured_estimate_blocks_the_overrun_the_static_floor_admits(monkeypatch): + """Concurrent unbounded requests must stop at the declared budget. + + A key with tpm_limit=8000 whose model really emits ~3000 output tokens + admits 7 concurrent requests under the 1024 floor (7 * 1025 <= 8000), so + once they all report actual usage the window carries ~21000 tokens + against an 8000 limit. Declaring the real output size admits only the two + requests the budget actually covers. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + + async def admitted(metadata): + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token(f"sk-overrun-{metadata}"), + tpm_limit=8000, + metadata=metadata, + ) + accepted = 0 + for _ in range(10): + try: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + call_type="completion", + ) + except HTTPException: + break + accepted += 1 + return accepted + + assert await admitted({}) == 7 + assert await admitted({"default_estimated_output_tokens": 3000}) == 2 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 69f04ce2bbe..2b162774aea 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -17,6 +17,7 @@ from litellm.proxy.hooks.proxy_track_cost_callback import ( _should_track_cost_callback, _update_database_and_spend_counters, ) +from litellm.types.utils import CallTypes @pytest.mark.asyncio @@ -783,6 +784,166 @@ async def test_enrich_failure_metadata_skips_when_no_api_key(): mock_get_key.assert_not_called() +@pytest.mark.asyncio +async def test_enrich_failure_metadata_keeps_captured_identity_when_not_resolving(): + """ + With resolve_missing_key_identity=False the key is not read, so a null user_id, + team_id and org_id captured earlier stay null instead of being refilled from the + key as it stands now. The team_alias lookup still runs off the captured team_id. + """ + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "alias-assigned-later" + mock_key_obj.user_id = "user-assigned-later" + mock_key_obj.team_id = "team-assigned-later" + mock_key_obj.org_id = "org-assigned-later" + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "captured-team-alias" + + with ( + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ) as mock_get_key, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ), + ): + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": None, + "user_api_key_user_id": None, + "user_api_key_team_id": "captured-team-id", + "user_api_key_team_alias": None, + "user_api_key_org_id": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( + metadata, resolve_missing_key_identity=False + ) + + mock_get_key.assert_not_called() + assert result["user_api_key_user_id"] is None + assert result["user_api_key_team_id"] == "captured-team-id" + assert result["user_api_key_org_id"] is None + assert result["user_api_key_alias"] is None + assert result["user_api_key_team_alias"] == "captured-team-alias" + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_ignores_flag_when_alias_present(): + """ + A captured alias already closes the key lookup, so resolve_missing_key_identity + changes nothing for a key that has one; only the alias-less key depends on it. + """ + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "captured-team-alias" + + with ( + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ), + ): + for resolve in (True, False): + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": "captured-alias", + "user_api_key_user_id": None, + "user_api_key_team_id": "captured-team-id", + "user_api_key_team_alias": None, + "user_api_key_org_id": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( + metadata, resolve_missing_key_identity=resolve + ) + mock_get_key.assert_not_called() + assert result["user_api_key_user_id"] is None + assert result["user_api_key_alias"] == "captured-alias" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type, expect_key_read", + [ + (CallTypes.aretrieve_batch.value, False), + (CallTypes.aretrieve_batch, False), + (CallTypes.acompletion.value, True), + ], +) +async def test_track_cost_callback_reads_key_only_for_in_request_logs(call_type, expect_key_read): + """ + The batch cost row is logged long after the batch was created, so it keeps the + identity persisted at create time. Every other call type still backfills from + the key. + """ + logger = _ProxyDBLogger() + + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "alias-assigned-later" + mock_key_obj.user_id = "user-assigned-later" + mock_key_obj.team_id = "team-assigned-later" + mock_key_obj.org_id = "org-assigned-later" + + kwargs = { + "call_type": call_type, + "model": None, + "litellm_call_id": "test-call-id", + "stream": False, + "litellm_params": { + "metadata": { + "user_api_key": "hashed_key", + "user_api_key_alias": None, + "user_api_key_user_id": None, + "user_api_key_team_id": None, + "user_api_key_org_id": None, + } + }, + } + + with ( + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ) as mock_get_key, + patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=MagicMock(team_alias=None), + ), + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging, + ): + mock_proxy_logging.failed_tracking_alert = AsyncMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + + await logger._PROXY_track_cost_callback( + kwargs=kwargs, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert mock_get_key.called is expect_key_read + + written = kwargs["litellm_params"]["metadata"] + if expect_key_read: + assert written["user_api_key_user_id"] == "user-assigned-later" + assert written["user_api_key_team_id"] == "team-assigned-later" + else: + assert written["user_api_key_user_id"] is None + assert written["user_api_key_team_id"] is None + assert written["user_api_key_org_id"] is None + + @pytest.mark.asyncio async def test_async_post_call_failure_hook_enriches_auth_error_metadata(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_common.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_common.py new file mode 100644 index 00000000000..f3515e84d0d --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_common.py @@ -0,0 +1,156 @@ +import ast +import importlib.util +from pathlib import Path +from types import ModuleType +from typing import Annotated + +import fastapi.dependencies.utils as fastapi_dependency_utils +import pytest +from fastapi import Depends, FastAPI, Header, Query, Request +from fastapi.testclient import TestClient + +import litellm.proxy.management_endpoints.management_v1.common as common_module +from litellm.proxy.management_endpoints.management_v1.common import ( + PROBLEM_CONTENT_TYPE, + ManagementProblem, + _declared_query_params, + problem_response, + reject_unknown_query_params, +) + + +def _client() -> TestClient: + app = FastAPI() + + @app.exception_handler(ManagementProblem) + async def _handle(_request: Request, exc: ManagementProblem): + return problem_response(exc.problem) + + @app.get("/things/{thing_id}", dependencies=[Depends(reject_unknown_query_params)]) + def _handler( + thing_id: str, + request: Request, + status: Annotated[str | None, Query(alias="filter[status]")] = None, + page: Annotated[int, Query(ge=1)] = 1, + x_trace: Annotated[str | None, Header()] = None, + ) -> dict[str, bool]: + return {"ok": True} + + return TestClient(app, raise_server_exceptions=False) + + +def test_a_declared_query_param_is_accepted_by_its_alias(): + response = _client().get("/things/abc", params={"filter[status]": "active", "page": "2"}) + assert response.status_code == 200, response.text + + +def test_an_unknown_query_param_is_rejected_as_a_problem(): + response = _client().get("/things/abc", params={"bogus": "x"}) + assert response.status_code == 400 + assert response.headers["content-type"].startswith(PROBLEM_CONTENT_TYPE) + assert "bogus" in response.json()["detail"] + + +def test_a_path_param_name_is_not_a_declared_query_param(): + """The flatten step returns path+query+header together; only query names count as declared. + + If the ParamTypes.query filter were dropped, `thing_id` (a path param) would leak + into the declared set and this request would be wrongly accepted. + """ + response = _client().get("/things/abc", params={"thing_id": "x"}) + assert response.status_code == 400 + assert "thing_id" in response.json()["detail"] + + +def test_a_header_param_name_is_not_a_declared_query_param(): + response = _client().get("/things/abc", params={"x-trace": "x"}) + assert response.status_code == 400 + assert "x-trace" in response.json()["detail"] + + +def test_declared_query_params_isolates_query_aliases_from_other_param_types(): + captured: dict[str, frozenset[str]] = {} + app = FastAPI() + + @app.get("/things/{thing_id}") + def _handler( + thing_id: str, + request: Request, + status: Annotated[str | None, Query(alias="filter[status]")] = None, + page: Annotated[int, Query(ge=1)] = 1, + x_trace: Annotated[str | None, Header()] = None, + ) -> dict[str, bool]: + captured["declared"] = _declared_query_params(request) + return {"ok": True} + + TestClient(app).get("/things/abc") + assert captured["declared"] == frozenset({"filter[status]", "page"}) + + +def test_declared_query_params_is_empty_when_the_route_has_no_dependant(): + request = Request( + { + "type": "http", + "method": "GET", + "scheme": "http", + "root_path": "", + "path": "/things/abc", + "query_string": b"", + "headers": [(b"host", b"testserver")], + } + ) + assert _declared_query_params(request) == frozenset() + + +# fastapi removed these in 0.140.7, which `pyproject.toml` still allows via +# `fastapi>=0.136.3,<1.0`. Add a name here whenever a supported release drops one. +FASTAPI_NAMES_REMOVED_IN_0_140_7 = frozenset({"get_flat_dependant"}) + +MANAGEMENT_V1_PACKAGE = Path(str(common_module.__file__)).parent + + +def _public_names(module: ModuleType) -> frozenset[str]: + return frozenset(name for name in vars(module) if not name.startswith("_")) + + +def _fastapi_names_imported_by(source_file: Path) -> frozenset[str]: + tree = ast.parse(source_file.read_text()) + return frozenset( + alias.name + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) and (node.module or "").startswith("fastapi") + for alias in node.names + ) + + +@pytest.mark.parametrize( + "source_file", sorted(MANAGEMENT_V1_PACKAGE.glob("*.py")), ids=lambda path: path.name +) +def test_no_module_imports_a_fastapi_name_removed_in_a_supported_release(source_file: Path): + """`pyproject.toml` allows fastapi up to <1.0, but CI only ever resolves 0.136.3. + + Every other test here passes just as well against a module importing a name + fastapi has since deleted, because the pinned fastapi still has it. On a user's + fastapi>=0.140.7 that import is an ImportError, and `proxy_server` imports this + package unguarded at module level, so it takes the whole proxy down rather than + just these routes. Globbing the package means a new module is covered on sight. + """ + assert not _fastapi_names_imported_by(source_file) & FASTAPI_NAMES_REMOVED_IN_0_140_7 + + +def test_common_still_imports_when_fastapi_has_dropped_those_names(monkeypatch: pytest.MonkeyPatch): + """The static check above cannot prove the module actually loads; this does. + + Behaviour cannot be asserted under the same simulation: on 0.136.3 + `get_flat_params` calls `get_flat_dependant` internally, so it raises NameError + once the name is gone. Loading is the part this pins. + """ + for name in FASTAPI_NAMES_REMOVED_IN_0_140_7: + monkeypatch.delattr(fastapi_dependency_utils, name, raising=False) + spec = importlib.util.spec_from_file_location( + "management_v1_common__simulated_fastapi", Path(str(common_module.__file__)) + ) + assert spec is not None and spec.loader is not None + reimported = importlib.util.module_from_spec(spec) + spec.loader.exec_module(reimported) + assert _public_names(reimported) == _public_names(common_module) diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index 3bdf9bafdc7..6a9e894feb5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -72,6 +72,39 @@ async def test_new_budget_success(client_and_mocks): mock_table.create.assert_awaited_once() +@pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) +@pytest.mark.asyncio +async def test_new_budget_rejects_a_duration_that_never_advances( + client_and_mocks, bad_duration +): + """A zero-length window resets to "now", so the row is due again the moment + it is written and the reset job re-reads it on every tick forever.""" + client, _, mock_table = client_and_mocks + + resp = client.post( + "/budget/new", + json={"budget_id": "budget_bad", "max_budget": 10.0, "budget_duration": bad_duration}, + ) + + assert resp.status_code == 400, resp.text + assert "Invalid budget_duration" in resp.json()["detail"]["error"] + mock_table.create.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_budget_rejects_a_duration_that_never_advances(client_and_mocks): + client, _, mock_table = client_and_mocks + + resp = client.post( + "/budget/update", + json={"budget_id": "budget_456", "budget_duration": "0s"}, + ) + + assert resp.status_code == 400, resp.text + assert "Invalid budget_duration" in resp.json()["detail"]["error"] + mock_table.update.assert_not_awaited() + + @pytest.mark.asyncio async def test_new_budget_db_not_connected(client_and_mocks, monkeypatch): client, mock_prisma, mock_table = client_and_mocks diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 469e0d340f0..ab9b4bc3922 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -7,9 +7,9 @@ from unittest.mock import AsyncMock, MagicMock import pytest -sys.path.insert( - 0, os.path.abspath("../../../..") -) # Adds the parent directory to the system path +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + +sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path from litellm.proxy.management_endpoints.common_daily_activity import ( _adjust_dates_for_timezone, @@ -108,8 +108,7 @@ async def test_get_daily_activity_order_has_id_tiebreaker(): mock_table.find_many.assert_called_once() order = mock_table.find_many.call_args[1]["order"] assert order == [{"date": "desc"}, {"id": "asc"}], ( - f"order must include the id tiebreaker after date for stable offset " - f"pagination (see #30164); got {order!r}" + f"order must include the id tiebreaker after date for stable offset pagination (see #30164); got {order!r}" ) @@ -301,9 +300,7 @@ async def test_get_api_key_metadata_returns_active_key_metadata(): mock_active_key.key_alias = "my-active-key" mock_active_key.team_id = "team-abc" - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[mock_active_key] - ) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[mock_active_key]) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -329,9 +326,7 @@ async def test_get_api_key_metadata_falls_back_to_deleted_keys(): mock_deleted_key.key_alias = "toto-test-2" mock_deleted_key.team_id = "team-xyz" - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( - return_value=[mock_deleted_key] - ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[mock_deleted_key]) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -360,9 +355,7 @@ async def test_get_api_key_metadata_mixed_active_and_deleted_keys(): mock_active_key.key_alias = "active-alias" mock_active_key.team_id = "team-active" - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[mock_active_key] - ) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[mock_active_key]) # One deleted key found mock_deleted_key = MagicMock() @@ -370,9 +363,7 @@ async def test_get_api_key_metadata_mixed_active_and_deleted_keys(): mock_deleted_key.key_alias = "deleted-alias" mock_deleted_key.team_id = "team-deleted" - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( - return_value=[mock_deleted_key] - ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[mock_deleted_key]) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -397,13 +388,9 @@ async def test_get_api_key_metadata_deleted_table_not_queried_when_all_keys_foun mock_active_key.key_alias = "alias-1" mock_active_key.team_id = "team-1" - mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock( - return_value=[mock_active_key] - ) + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[mock_active_key]) mock_prisma.db.litellm_deletedverificationtoken = MagicMock() - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( - return_value=[] - ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[]) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -425,9 +412,7 @@ async def test_get_api_key_metadata_deleted_table_error_handled_gracefully(): mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) # Deleted table raises an error (e.g., table doesn't exist in older schema) - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( - side_effect=Exception("Table not found") - ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(side_effect=Exception("Table not found")) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -458,9 +443,7 @@ async def test_get_api_key_metadata_regenerated_key_uses_most_recent_deleted_rec mock_deleted_2.team_id = "older-team" # Ordered by deleted_at desc, so first record is the most recent - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( - return_value=[mock_deleted_1, mock_deleted_2] - ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[mock_deleted_1, mock_deleted_2]) result = await get_api_key_metadata( prisma_client=mock_prisma, @@ -633,9 +616,7 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): mock_deleted_key.team_id = "69cd4b77-b095-4489-8c46-4f2f31d840a2" mock_prisma.db.litellm_deletedverificationtoken = MagicMock() - mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock( - return_value=[mock_deleted_key] - ) + mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[mock_deleted_key]) result = await get_daily_activity_aggregated( prisma_client=mock_prisma, @@ -754,15 +735,9 @@ async def test_model_groups_breakdown_keys_by_public_name_with_model_fallback(): mock_prisma.db = MagicMock() records = [ - _daily_user_spend_record( - user_id="u1", api_key="key-1", spend=7.0, model="gpt-5.2", model_group="gpt-5.2-eu" - ), - _daily_user_spend_record( - user_id="u1", api_key="key-1", spend=3.0, model="gpt-5.2", model_group=None - ), - _daily_user_spend_record( - user_id="u1", api_key="key-1", spend=2.0, model="claude-x", model_group="" - ), + _daily_user_spend_record(user_id="u1", api_key="key-1", spend=7.0, model="gpt-5.2", model_group="gpt-5.2-eu"), + _daily_user_spend_record(user_id="u1", api_key="key-1", spend=3.0, model="gpt-5.2", model_group=None), + _daily_user_spend_record(user_id="u1", api_key="key-1", spend=2.0, model="claude-x", model_group=""), ] mock_table = MagicMock() @@ -829,9 +804,7 @@ class TestAdjustDatesForTimezone: ], ) def test_returns_input_dates_unchanged_for_any_offset(self, offset_minutes): - start, end = _adjust_dates_for_timezone( - "2026-05-29", "2026-05-29", offset_minutes - ) + start, end = _adjust_dates_for_timezone("2026-05-29", "2026-05-29", offset_minutes) assert start == "2026-05-29" assert end == "2026-05-29" @@ -859,9 +832,7 @@ class TestAdjustDatesForTimezone: exceeded the multi-day total by ~50% over a 5-day IST window. """ days = ["2026-05-29", "2026-05-30", "2026-05-31", "2026-06-01", "2026-06-02"] - single_day_ranges = [ - _adjust_dates_for_timezone(d, d, offset_minutes) for d in days - ] + single_day_ranges = [_adjust_dates_for_timezone(d, d, offset_minutes) for d in days] multi_day_range = _adjust_dates_for_timezone(days[0], days[-1], offset_minutes) per_day_starts = [r[0] for r in single_day_ranges] @@ -894,9 +865,7 @@ class TestAdjustDatesForTimezoneLiveEnd: assert (start, end) == ("2026-07-06", "2026-08-06") def test_without_opt_in_live_range_keeps_pass_through(self): - start, end = _adjust_dates_for_timezone( - "2026-07-06", "2026-08-05", 420, utc_now=self.PT_EVENING_UTC - ) + start, end = _adjust_dates_for_timezone("2026-07-06", "2026-08-05", 420, utc_now=self.PT_EVENING_UTC) assert (start, end) == ("2026-07-06", "2026-08-05") def test_pt_historical_range_is_untouched(self): @@ -1173,3 +1142,580 @@ class TestEverySavingsDriverSurvivesTheReadPath: assert f"total_{driver}" in DailySpendMetadata.model_fields, ( f"total_{driver} is missing, so the range summary omits the driver" ) + + +@pytest.fixture +def ptu_cost_attribution_enabled(monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + +def _spend_record(api_key, *, model="gpt-4o-mini-ptu", spend=0.0, ptu_flat_cost=0.0): + return SimpleNamespace( + api_key=api_key, + model=model, + model_group=None, + mcp_namespaced_tool_name=None, + custom_llm_provider="openai", + endpoint=None, + spend=spend, + prompt_tokens=0, + completion_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + compression_saved_tokens=0, + compression_savings_spend=0, + prompt_caching_savings_spend=0, + autorouter_savings_spend=0, + total_tokens=0, + api_requests=0, + successful_requests=0, + failed_requests=0, + ptu_flat_cost=ptu_flat_cost, + ) + + +def test_update_metrics_accumulates_ptu_flat_cost(ptu_cost_attribution_enabled): + metrics = update_metrics(SpendMetrics(), _spend_record("real-key", spend=1.0, ptu_flat_cost=240.0)) + assert metrics.flat_cost == 240.0 + assert metrics.spend == 1.0 + + +def test_ptu_sentinel_excluded_from_key_breakdown_but_flat_cost_aggregates(ptu_cost_attribution_enabled): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + update_breakdown_metrics(breakdown, _spend_record("real-key", spend=5.0, ptu_flat_cost=0.0), {}, {}, {}) + update_breakdown_metrics(breakdown, _spend_record(PTU_SENTINEL_API_KEY, spend=0.0, ptu_flat_cost=240.0), {}, {}, {}) + + model_bucket = breakdown.models["gpt-4o-mini-ptu"] + # flat cost aggregates into the parent model metrics + assert model_bucket.metrics.flat_cost == 240.0 + assert model_bucket.metrics.spend == 5.0 + # the sentinel never appears as an api_key row; only the real key does + assert PTU_SENTINEL_API_KEY not in model_bucket.api_key_breakdown + assert "real-key" in model_bucket.api_key_breakdown + + +def _grouping_row( + group_level, + *, + api_key=None, + model=None, + model_group=None, + custom_llm_provider="openai", + mcp_namespaced_tool_name=None, + endpoint=None, + spend=0.0, + ptu_flat_cost=0.0, +): + from litellm.proxy.management_endpoints.common_daily_activity import _GroupingSetsRow + + return _GroupingSetsRow( + date="2024-01-01", + api_key=api_key, + model=model, + model_group=model_group, + custom_llm_provider=custom_llm_provider, + mcp_namespaced_tool_name=mcp_namespaced_tool_name, + endpoint=endpoint, + group_level=group_level, + spend=spend, + ptu_flat_cost=ptu_flat_cost, + prompt_tokens=0, + completion_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + compression_saved_tokens=0, + compression_savings_spend=0.0, + prompt_caching_savings_spend=0.0, + autorouter_savings_spend=0.0, + api_requests=0, + successful_requests=0, + failed_requests=0, + ) + + +def test_grouping_sets_dispatcher_excludes_ptu_sentinel_from_key_breakdowns(ptu_cost_attribution_enabled): + """The GROUPING SETS path must mirror the per-row path: the flat-cost sentinel + aggregates into the date/model/total metrics but never surfaces as an api_key.""" + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_API_KEY, + _GROUP_DATE_MODEL, + _GROUP_DATE_MODEL_API_KEY, + _GROUP_GRAND_TOTAL, + _aggregate_grouping_sets_records_sync, + ) + + records = [ + _grouping_row(_GROUP_DATE_API_KEY, api_key="real-key", spend=5.0), + _grouping_row(_GROUP_DATE_API_KEY, api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + _grouping_row(_GROUP_DATE_MODEL, model="gpt-4o-mini-ptu", spend=5.0, ptu_flat_cost=240.0), + _grouping_row(_GROUP_DATE_MODEL_API_KEY, model="gpt-4o-mini-ptu", api_key="real-key", spend=5.0), + _grouping_row( + _GROUP_DATE_MODEL_API_KEY, model="gpt-4o-mini-ptu", api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0 + ), + _grouping_row(_GROUP_GRAND_TOTAL, spend=5.0, ptu_flat_cost=240.0), + ] + + aggregated = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={}) + + assert aggregated["totals"].flat_cost == 240.0 + day = aggregated["results"][0] + assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys + assert "real-key" in day.breakdown.api_keys + + model_bucket = day.breakdown.models["gpt-4o-mini-ptu"] + assert model_bucket.metrics.flat_cost == 240.0 + assert model_bucket.metrics.spend == 5.0 + assert PTU_SENTINEL_API_KEY not in model_bucket.api_key_breakdown + assert "real-key" in model_bucket.api_key_breakdown + + +def test_grouping_sets_dispatcher_populates_every_breakdown_level(ptu_cost_attribution_enabled): + """Every GROUPING SETS level lands in its bucket, and the flat-cost sentinel + is kept out of the model_group and provider api_key sub-breakdowns too.""" + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_ENDPOINT, + _GROUP_DATE_ENDPOINT_API_KEY, + _GROUP_DATE_MCP, + _GROUP_DATE_MCP_API_KEY, + _GROUP_DATE_MODEL_GROUP, + _GROUP_DATE_MODEL_GROUP_API_KEY, + _GROUP_DATE_PROVIDER, + _GROUP_DATE_PROVIDER_API_KEY, + _aggregate_grouping_sets_records_sync, + ) + + records = [ + _grouping_row(_GROUP_DATE_MODEL_GROUP, model_group="grp", spend=4.0, ptu_flat_cost=240.0), + _grouping_row(_GROUP_DATE_MODEL_GROUP_API_KEY, model_group="grp", api_key="real-key", spend=4.0), + _grouping_row( + _GROUP_DATE_MODEL_GROUP_API_KEY, model_group="grp", api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0 + ), + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="azure", spend=4.0), + _grouping_row(_GROUP_DATE_PROVIDER_API_KEY, custom_llm_provider="azure", api_key="real-key", spend=4.0), + _grouping_row( + _GROUP_DATE_PROVIDER_API_KEY, + custom_llm_provider="azure", + api_key=PTU_SENTINEL_API_KEY, + ptu_flat_cost=240.0, + ), + _grouping_row(_GROUP_DATE_MCP, mcp_namespaced_tool_name="srv/tool", spend=2.0), + _grouping_row(_GROUP_DATE_MCP_API_KEY, mcp_namespaced_tool_name="srv/tool", api_key="real-key", spend=2.0), + _grouping_row(_GROUP_DATE_ENDPOINT, endpoint="/v1/chat/completions", spend=3.0), + _grouping_row(_GROUP_DATE_ENDPOINT_API_KEY, endpoint="/v1/chat/completions", api_key="real-key", spend=3.0), + ] + + aggregated = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={}) + day = aggregated["results"][0] + + group_bucket = day.breakdown.model_groups["grp"] + assert group_bucket.metrics.flat_cost == 240.0 + assert PTU_SENTINEL_API_KEY not in group_bucket.api_key_breakdown + assert "real-key" in group_bucket.api_key_breakdown + + provider_bucket = day.breakdown.providers["azure"] + assert PTU_SENTINEL_API_KEY not in provider_bucket.api_key_breakdown + assert "real-key" in provider_bucket.api_key_breakdown + + assert "real-key" in day.breakdown.mcp_servers["srv/tool"].api_key_breakdown + assert "real-key" in day.breakdown.endpoints["/v1/chat/completions"].api_key_breakdown + + +def test_grouping_sets_dispatcher_keeps_ptu_flat_cost_out_of_the_provider_breakdown(): + """Sentinel rows carry no provider, so their flat cost must not surface under the + "unknown" provider - the per-row path skips them for exactly the same reason.""" + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_PROVIDER, + _aggregate_grouping_sets_records_sync, + ) + + records = [ + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="azure", spend=4.0), + # the sentinel's own provider-level row: empty provider, flat cost only + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="", ptu_flat_cost=240.0), + ] + + aggregated = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={}) + providers = aggregated["results"][0].breakdown.providers + + # the bucket is still reported (a legacy all-zero row must not vanish); only the + # flat cost is withheld, so no provider is credited with PTU capacity cost + assert providers["azure"].metrics.spend == 4.0 + assert sum(bucket.metrics.flat_cost for bucket in providers.values()) == 0.0 + + +def test_grouping_sets_dispatcher_keeps_a_real_provider_row_that_shares_the_sentinel_shape(): + """A request row whose provider is empty still gets its "unknown" bucket - only the + flat cost is withheld, so provider attribution of real spend is unchanged.""" + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_PROVIDER, + _aggregate_grouping_sets_records_sync, + ) + + records = [_grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="", spend=4.0, ptu_flat_cost=240.0)] + + aggregated = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={}) + unknown = aggregated["results"][0].breakdown.providers["unknown"] + + assert unknown.metrics.spend == 4.0 + assert unknown.metrics.flat_cost == 0.0 + + +def test_update_breakdown_metrics_covers_mcp_endpoint_and_entity(ptu_cost_attribution_enabled): + """A full request record fans out into the mcp, endpoint, provider and entity + breakdowns, while the flat-cost sentinel stays out of the entity api_key sub-map.""" + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + record = SimpleNamespace( + api_key="real-key", + model="gpt-4o-mini-ptu", + model_group="grp", + mcp_namespaced_tool_name="srv/tool", + custom_llm_provider="azure", + endpoint="/v1/chat/completions", + spend=5.0, + prompt_tokens=0, + completion_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + compression_saved_tokens=0, + compression_savings_spend=0, + prompt_caching_savings_spend=0, + autorouter_savings_spend=0, + total_tokens=0, + api_requests=0, + successful_requests=0, + failed_requests=0, + ptu_flat_cost=0.0, + team_id="team-1", + ) + update_breakdown_metrics(breakdown, record, {}, {}, {}, entity_id_field="team_id") + + assert "srv/tool" in breakdown.mcp_servers + assert "real-key" in breakdown.mcp_servers["srv/tool"].api_key_breakdown + assert "/v1/chat/completions" in breakdown.endpoints + assert "azure" in breakdown.providers + assert "team-1" in breakdown.entities + assert "real-key" in breakdown.entities["team-1"].api_key_breakdown + + sentinel = SimpleNamespace(**{**record.__dict__, "api_key": PTU_SENTINEL_API_KEY, "ptu_flat_cost": 240.0}) + update_breakdown_metrics(breakdown, sentinel, {}, {}, {}, entity_id_field="team_id") + assert PTU_SENTINEL_API_KEY not in breakdown.entities["team-1"].api_key_breakdown + assert breakdown.entities["team-1"].metrics.flat_cost == 240.0 + + +def test_grouping_sets_dispatcher_keeps_an_all_zero_legacy_provider_bucket(): + """LiteLLM_DailyTeamSpend predates its api_requests column; the migration that added it + backfilled NOT NULL DEFAULT 0, so a legacy keyless row is all zeroes. Dropping those + would silently remove a provider the base build reported.""" + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_PROVIDER, + _aggregate_grouping_sets_records_sync, + ) + + records = [ + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="ollama"), # spend/tokens/requests all 0 + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="openai", spend=0.25), + ] + + providers = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={})["results"][ + 0 + ].breakdown.providers + + assert set(providers) == {"ollama", "openai"} + assert providers["ollama"].metrics.spend == 0.0 + assert providers["ollama"].metrics.flat_cost == 0.0 + + +class TestSentinelRowsDisplayTheirModelName: + """A sentinel row keys on the deployment id so a rename cannot move it. The usage views + render the breakdown key directly as a label, so the read path has to show the name.""" + + @pytest.fixture(autouse=True) + def _enabled(self, ptu_cost_attribution_enabled): + """Flat cost is gated off by default, and these assert on the amounts.""" + + @staticmethod + def _breakdown(records): + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + for record in records: + update_breakdown_metrics(breakdown, record, {}, {}, {}) + return breakdown + + @staticmethod + def _sentinel(*, model_id, model_group, flat_cost=480.0): + from litellm.constants import PTU_SENTINEL_API_KEY + + record = _spend_record(PTU_SENTINEL_API_KEY, model=model_id, spend=0.0, ptu_flat_cost=flat_cost) + record.model_group = model_group + return record + + def test_models_breakdown_keys_a_sentinel_row_on_its_public_name(self): + models = self._breakdown([self._sentinel(model_id="dep-1", model_group="gpt-4o-ptu")]).models + + assert "gpt-4o-ptu" in models, f"the UI would label this row a UUID: {list(models)}" + assert "dep-1" not in models + assert models["gpt-4o-ptu"].metrics.flat_cost == pytest.approx(480.0) + + def test_two_deployments_sharing_a_name_merge_under_it(self): + """The write path stopped collapsing them, so the read path has to.""" + models = self._breakdown( + [ + self._sentinel(model_id="dep-a", model_group="gpt-4o-ptu", flat_cost=240.0), + self._sentinel(model_id="dep-b", model_group="gpt-4o-ptu", flat_cost=120.0), + ] + ).models + + assert list(models) == ["gpt-4o-ptu"] + assert models["gpt-4o-ptu"].metrics.flat_cost == pytest.approx(360.0) + + def test_a_request_row_still_keys_on_its_model(self): + """Scoped to sentinel rows: a request row keys on model as it always has, even + though it also carries a model_group.""" + record = _spend_record("real-key", model="gemini/gemini-2.5-flash", spend=1.25) + record.model_group = "gemini-live" + + models = self._breakdown([record]).models + + assert "gemini/gemini-2.5-flash" in models + assert "gemini-live" not in models + + def test_a_sentinel_row_without_a_model_group_falls_back_to_the_id(self): + """Never drop the charge: an unexpected row with no display name still reports.""" + models = self._breakdown([self._sentinel(model_id="dep-1", model_group=None)]).models + + assert models["dep-1"].metrics.flat_cost == pytest.approx(480.0) + + +def _daily_team_row(api_key, *, spend=0.0, ptu_flat_cost=0.0): + """A LiteLLM_DailyTeamSpend row as the paginated read path receives it from find_many.""" + base: Final = _spend_record(api_key, spend=spend, ptu_flat_cost=ptu_flat_cost) + return SimpleNamespace(**{**base.__dict__, "date": "2026-07-01", "team_id": "team-1"}) + + +class TestPtuCostAttributionDisabled: + """With LITELLM_ENABLE_PTU_COST_ATTRIBUTION unset, both read paths report zero flat + cost, while the sentinel filtering that keeps ``__ptu_flat_cost__`` out of the + breakdowns keeps running. + + Filtering is deliberately not gated: an operator can enable the flag, accrue + sentinel rows, then disable it, and those rows stay in LiteLLM_DailyTeamSpend + forever. Gating the filter too would surface the sentinel as a bogus api_key and + mint a provider bucket for its empty provider. + """ + + @pytest.fixture(autouse=True) + def _flag_off(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + def test_paginated_path_reports_zero_flat_cost(self): + metrics = update_metrics(SpendMetrics(), _spend_record("real-key", spend=1.0, ptu_flat_cost=240.0)) + + assert metrics.flat_cost == 0.0 + assert metrics.spend == 1.0 + + def test_aggregated_path_reports_zero_flat_cost(self): + from litellm.proxy.management_endpoints.common_daily_activity import _GROUP_GRAND_TOTAL + + metrics = _record_to_spend_metrics(_grouping_row(_GROUP_GRAND_TOTAL, spend=5.0, ptu_flat_cost=240.0)) + + assert metrics.flat_cost == 0.0 + assert metrics.spend == 5.0 + + def test_aggregated_totals_and_buckets_report_zero_flat_cost(self): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_API_KEY, + _GROUP_DATE_MODEL, + _GROUP_GRAND_TOTAL, + _aggregate_grouping_sets_records_sync, + ) + + records = [ + _grouping_row(_GROUP_DATE_API_KEY, api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + _grouping_row(_GROUP_DATE_MODEL, model="gpt-4o-mini-ptu", spend=5.0, ptu_flat_cost=240.0), + _grouping_row(_GROUP_GRAND_TOTAL, spend=5.0, ptu_flat_cost=240.0), + ] + + aggregated = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={}) + + assert aggregated["totals"].flat_cost == 0.0 + assert aggregated["totals"].spend == 5.0 + assert aggregated["results"][0].breakdown.models["gpt-4o-mini-ptu"].metrics.flat_cost == 0.0 + + def test_sentinel_still_excluded_from_the_api_key_breakdown(self): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + update_breakdown_metrics(breakdown, _spend_record("real-key", spend=5.0), {}, {}, {}) + update_breakdown_metrics( + breakdown, _spend_record(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), {}, {}, {}, entity_id_field="team_id" + ) + + assert PTU_SENTINEL_API_KEY not in breakdown.api_keys + assert PTU_SENTINEL_API_KEY not in breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown + assert "real-key" in breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown + + def test_sentinel_still_excluded_from_the_provider_breakdown(self): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + update_breakdown_metrics(breakdown, _spend_record(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), {}, {}, {}) + + assert breakdown.providers == {} + + def test_grouping_sets_sentinel_still_excluded_from_breakdowns(self): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_API_KEY, + _GROUP_DATE_MODEL, + _GROUP_DATE_MODEL_API_KEY, + _GROUP_DATE_PROVIDER, + _aggregate_grouping_sets_records_sync, + ) + + records = [ + _grouping_row(_GROUP_DATE_API_KEY, api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + _grouping_row(_GROUP_DATE_MODEL, model="gpt-4o-mini-ptu", spend=5.0, ptu_flat_cost=240.0), + _grouping_row( + _GROUP_DATE_MODEL_API_KEY, model="gpt-4o-mini-ptu", api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0 + ), + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="", ptu_flat_cost=240.0), + ] + + day = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={})["results"][0] + + assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys + assert PTU_SENTINEL_API_KEY not in day.breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown + assert sum(bucket.metrics.flat_cost for bucket in day.breakdown.providers.values()) == 0.0 + + @pytest.mark.asyncio + async def test_team_daily_activity_endpoint_reports_zero_flat_cost(self): + """/team/daily/activity reads rows with find_many rather than the aggregated SQL, so + forcing the SQL select to a constant zero would leave this path reporting flat cost.""" + from litellm.constants import PTU_SENTINEL_API_KEY + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_table = MagicMock() + mock_table.count = AsyncMock(return_value=2) + mock_table.find_many = AsyncMock( + return_value=[ + _daily_team_row("real-key", spend=5.0), + _daily_team_row(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + ] + ) + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_dailyteamspend = mock_table + + result = await get_daily_activity( + prisma_client=mock_prisma, + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-07-01", + end_date="2026-07-01", + model=None, + api_key=None, + page=1, + page_size=50, + ) + + assert result.metadata.total_flat_cost == 0.0 + assert result.metadata.total_spend == 5.0 + assert PTU_SENTINEL_API_KEY not in result.results[0].breakdown.api_keys + + @pytest.mark.asyncio + async def test_team_daily_activity_endpoint_reports_flat_cost_once_enabled(self, monkeypatch): + from litellm.constants import PTU_SENTINEL_API_KEY + + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_table = MagicMock() + mock_table.count = AsyncMock(return_value=2) + mock_table.find_many = AsyncMock( + return_value=[ + _daily_team_row("real-key", spend=5.0), + _daily_team_row(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + ] + ) + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_dailyteamspend = mock_table + + result = await get_daily_activity( + prisma_client=mock_prisma, + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-07-01", + end_date="2026-07-01", + model=None, + api_key=None, + page=1, + page_size=50, + ) + + assert result.metadata.total_flat_cost == 240.0 + assert PTU_SENTINEL_API_KEY not in result.results[0].breakdown.api_keys + + +class TestFlagIsNotReadOnTheHotPath: + """update_metrics runs once per accumulation and a record fans out across roughly a + dozen breakdowns, so a flag that reads through the secret manager must not be consulted + for rows that carry no flat cost at all.""" + + @staticmethod + def _count_flag_reads(records): + import litellm.proxy.management_endpoints.common_daily_activity as cda + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + reads = [] + real = cda.is_ptu_cost_attribution_enabled + + def counted(): + reads.append(1) + return real() + + cda.is_ptu_cost_attribution_enabled = counted + try: + breakdown = BreakdownMetrics() + for record in records: + cda.update_breakdown_metrics(breakdown, record, {}, {}, {}) + finally: + cda.is_ptu_cost_attribution_enabled = real + return len(reads) + + def test_a_request_row_never_reads_the_flag(self): + reads = self._count_flag_reads([_spend_record("real-key", spend=5.0, ptu_flat_cost=0.0)]) + assert reads == 0, f"{reads} secret-manager lookups for a row with no flat cost" + + def test_a_page_of_request_rows_never_reads_the_flag(self): + rows = [_spend_record(f"key-{i}", spend=1.0, ptu_flat_cost=0.0) for i in range(50)] + assert self._count_flag_reads(rows) == 0 + + def test_a_sentinel_row_still_consults_the_flag(self): + from litellm.constants import PTU_SENTINEL_API_KEY + + reads = self._count_flag_reads([_spend_record(PTU_SENTINEL_API_KEY, spend=0.0, ptu_flat_cost=240.0)]) + assert reads > 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index 81840745d0e..7dfd99dfa53 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -628,6 +628,58 @@ class TestValidateFiniteSpendErrorDetail: } +class TestValidateBudgetDuration: + """`validate_budget_duration` keeps durations that never advance out of the + database. + + A duration of "0s" resolves to a reset time of now, so the row is due again + the instant it is written. The reset job re-reads such rows on every tick + and, once one tenant owns enough of them, they fill each batch and starve + every other tenant's reset. + """ + + def test_none_is_allowed(self): + from litellm.proxy.management_endpoints.common_utils import ( + validate_budget_duration, + ) + + assert validate_budget_duration(None) is None + + @pytest.mark.parametrize("duration", ["30s", "5m", "1h", "1d", "7d", "30d", "1mo"]) + def test_positive_durations_are_allowed(self, duration): + from litellm.proxy.management_endpoints.common_utils import ( + validate_budget_duration, + ) + + assert validate_budget_duration(duration) is None + + @pytest.mark.parametrize("duration", ["0s", "0m", "0h", "0d", "-5m", "abc", ""]) + def test_non_advancing_durations_are_rejected(self, duration): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + validate_budget_duration, + ) + + with pytest.raises(HTTPException) as exc_info: + validate_budget_duration(duration) + assert exc_info.value.status_code == 400 + + def test_rejection_detail_is_exact(self): + from fastapi import HTTPException + + from litellm.proxy.management_endpoints.common_utils import ( + validate_budget_duration, + ) + + with pytest.raises(HTTPException) as exc_info: + validate_budget_duration("0s") + + assert exc_info.value.detail == { + "error": "Invalid budget_duration '0s'. Use a format like '1h', '24h', '7d', or '30d'." + } + + class TestRequireCallerUserIdErrorDetail: """The 403 for a service-account key must carry the exact error body.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 0af5ad6cd9b..5efed8de325 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -749,6 +749,40 @@ def test_char_new_body(mock_prisma_client, mock_user_api_key_auth): assert response.json() == _EXPECTED_CUSTOMER +@pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) +def test_customer_new_rejects_a_duration_that_never_advances( + mock_prisma_client, mock_user_api_key_auth, bad_duration +): + """A zero-length window resets to "now", leaving the customer's budget row + permanently due for the reset job to re-read every tick.""" + mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW)) + + response = client.post( + "/customer/new", + json={"user_id": "c1", "max_budget": 10.0, "budget_duration": bad_duration}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 400, response.text + assert "Invalid budget_duration" in response.text + mock_prisma_client.db.litellm_endusertable.create.assert_not_awaited() + + +def test_customer_new_accepts_a_normal_duration(mock_prisma_client, mock_user_api_key_auth): + mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW)) + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=_row({"budget_id": "b1", "max_budget": 10.0}) + ) + + response = client.post( + "/customer/new", + json={"user_id": "c1", "max_budget": 10.0, "budget_duration": "30d"}, + headers={"Authorization": "Bearer k"}, + ) + + assert response.status_code == 200, response.text + + def test_char_update_body(mock_prisma_client, mock_user_api_key_auth): mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( return_value=_row({"user_id": "c1", "blocked": False}) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 056c2d3657a..cd5a5d42b09 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -788,6 +788,68 @@ def test_update_internal_user_params_reset_spend_and_max_budget(): assert "budget_duration" not in non_default_values # Should not add default values +@pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) +def test_update_internal_user_params_rejects_a_duration_that_never_advances(bad_duration): + """A zero-length window resets to "now", so the user row is due again the + moment it is written and the reset job re-reads it on every tick. Enough of + them fill each batch and starve other tenants' resets. + """ + from fastapi import HTTPException + + from litellm.proxy._types import UpdateUserRequest + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_internal_user_params, + ) + + data = UpdateUserRequest(user_id="test_user_id", budget_duration=bad_duration) + + with pytest.raises(HTTPException) as exc_info: + _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data) + + assert exc_info.value.status_code == 400 + assert "Invalid budget_duration" in str(exc_info.value.detail) + + +def test_update_internal_user_params_accepts_a_normal_duration(): + from litellm.proxy._types import UpdateUserRequest + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_internal_user_params, + ) + + data = UpdateUserRequest(user_id="test_user_id", budget_duration="30d") + + non_default_values = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data) + + assert non_default_values["budget_duration"] == "30d" + assert non_default_values["budget_reset_at"] is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) +async def test_new_user_rejects_a_duration_that_never_advances(mocker, bad_duration): + """/user/new must reject the same never-advancing durations /user/update does.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.management_endpoints.internal_user_endpoints import new_user + + mocker.patch("litellm.proxy.proxy_server.prisma_client", MagicMock()) + duplicate_check = mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id", + new=AsyncMock(), + ) + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with pytest.raises(ProxyException) as exc_info: + await new_user( + data=NewUserRequest(budget_duration=bad_duration), + user_api_key_dict=admin, + ) + + assert str(exc_info.value.code) == "400" + assert "Invalid budget_duration" in str(exc_info.value.message) + duplicate_check.assert_not_awaited() + + @pytest.mark.asyncio async def test_new_user_license_over_limit(mocker): """ 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 e8709f3af34..8f151ed882c 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 @@ -2527,6 +2527,72 @@ def _setup_update_key_mocks(monkeypatch, mock_prisma_client): monkeypatch.setattr("litellm.store_audit_logs", False) +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) +async def test_update_key_rejects_a_duration_that_never_advances(monkeypatch, bad_duration): + """A zero-length window resets to "now", so the key row is due again the + moment it is written. The reset job re-reads such rows on every tick, and a + tenant with enough of them fills each batch and starves other tenants. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + hashed_token = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b" + key_in_db = LiteLLM_VerificationToken(token=hashed_token, user_id="test-user") + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + mock_prisma_client.update_data = AsyncMock() + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + with pytest.raises(ProxyException) as exc_info: + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=hashed_token, budget_duration=bad_duration), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + assert str(exc_info.value.code) == "400" + assert "Invalid budget_duration" in str(exc_info.value.message) + mock_prisma_client.update_data.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) +async def test_generate_key_rejects_a_duration_that_never_advances(monkeypatch, bad_duration): + """/key/generate must reject the same never-advancing durations /key/update does.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_fn, + ) + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new=AsyncMock(), + ) as mock_generate: + with pytest.raises(ProxyException) as exc_info: + await generate_key_fn( + data=GenerateKeyRequest(budget_duration=bad_duration), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234", user_id="1234" + ), + ) + + assert str(exc_info.value.code) == "400" + assert "Invalid budget_duration" in str(exc_info.value.message) + mock_generate.assert_not_awaited() + + @pytest.mark.asyncio async def test_update_key_by_alias_only(monkeypatch): """ @@ -8081,7 +8147,7 @@ async def test_key_with_budget_id_does_not_store_budget_duration(): budget_duration, the key does NOT get budget_duration stored on it. Keys with budget_id follow their linked budget tier's reset schedule; - reset_budget_for_keys_linked_to_budgets() resets them when the tier resets. + reset_budget_for_litellm_budget_table() resets them when the tier resets. This avoids duplicating budget_duration on keys so tier updates apply automatically to all linked keys. """ @@ -15442,3 +15508,312 @@ async def test_migrate_encryption_endpoint_rejects_proxy_admin_viewer(): assert exc_info.value.status_code == 403 mock_migrate.assert_not_awaited() + + +_ESTIMATE = "default_estimated_output_tokens" +_ESTIMATE_PER_MODEL = "default_estimated_output_tokens_per_model" + + +@pytest.mark.parametrize( + "label, request_body, existing_metadata, allowed", + [ + ("nothing declared", {}, None, True), + ("declared top-level on a key with none stored", {_ESTIMATE: 1}, None, False), + ("declared inside metadata on a key with none stored", {"metadata": {_ESTIMATE: 1}}, None, False), + ( + "per-model map declared inside metadata", + {"metadata": {_ESTIMATE_PER_MODEL: {"gpt-4": 1}}}, + None, + False, + ), + ("unrelated edit, metadata omitted", {"models": ["gpt-4"]}, {_ESTIMATE: 2000}, True), + ("stored value resent unchanged", {_ESTIMATE: 2000}, {_ESTIMATE: 2000}, True), + ("stored value lowered", {_ESTIMATE: 1}, {_ESTIMATE: 2000}, False), + ("stored value raised", {_ESTIMATE: 9000}, {_ESTIMATE: 2000}, False), + ( + "stored value cleared by sending a metadata blob without it", + {"metadata": {"other": "keep"}}, + {_ESTIMATE: 2000, "other": "keep"}, + False, + ), + ( + "stored value resent inside the metadata blob", + {"metadata": {_ESTIMATE: 2000, "other": "keep"}}, + {_ESTIMATE: 2000, "other": "keep"}, + True, + ), + ( + "per-model map resent unchanged", + {_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + {_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + True, + ), + ( + "one model in the per-model map lowered", + {_ESTIMATE_PER_MODEL: {"gpt-4": 1}}, + {_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + False, + ), + ], +) +def test_output_token_estimate_admin_gate_matrix(label, request_body, existing_metadata, allowed): + """A non-admin may only leave a key's stored output-token estimate exactly as it is. + + The estimate decides what the TPM limiter reserves for a request that omits + max_tokens, so lowering, raising or clearing it moves a reservation charged + against team and organization windows the key holder does not own. Key + metadata is writable by the key holder, and the declaration can be written + either as a dedicated top-level field or nested in the metadata blob, so + both routes are gated. Resending the stored value is what the edit form + produces on every save and has to stay allowed. + """ + from litellm.proxy.auth.auth_utils import ( + enforce_output_token_estimates_are_admin_only, + ) + + def _call(caller): + enforce_output_token_estimates_are_admin_only( + data=UpdateKeyRequest(key="sk-1", **request_body), + existing_metadata=existing_metadata, + user_api_key_dict=caller, + entity="key", + ) + + non_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-non-admin", + user_id="alice", + ) + if allowed: + _call(non_admin) + else: + with pytest.raises(HTTPException) as exc: + _call(non_admin) + assert exc.value.status_code == 403 + assert "Only proxy admins can set" in str(exc.value.detail) + + _call( + UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin", + ) + ) + + +@pytest.mark.asyncio +async def test_generate_key_output_token_estimate_rejected_for_non_admin(): + """The /key/update gate does not cover generate, so without its own check a + non-admin could self-mint a key that reserves one output token per + unbounded request and overrun the TPM window it is charged against.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(default_estimated_output_tokens=1, tpm_limit=100000), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "Only proxy admins can set" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_generate_key_output_token_estimate_in_metadata_rejected_for_non_admin(): + """Writing the declaration into the raw metadata blob lands in the same + stored field, so gating only the dedicated top-level field leaves the + bypass wide open.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(metadata={"default_estimated_output_tokens": 1}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + + +@pytest.mark.asyncio +async def test_generate_key_output_token_estimate_allowed_for_admin(): + """A proxy admin declaring the estimate must reach key creation.""" + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", False), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = { + "key": "sk-test-key", + "expires": None, + "user_id": "admin", + "team_id": None, + } + await _common_key_generation_helper( + data=GenerateKeyRequest(default_estimated_output_tokens=200), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"), + litellm_changed_by=None, + team_table=None, + ) + assert mock_generate_key.called + + +def _estimate_key_row(token: str, metadata: dict): + existing_key = MagicMock() + existing_key.token = token + existing_key.user_id = "internal_user" + existing_key.created_by = "internal_user" + existing_key.team_id = None + existing_key.project_id = None + existing_key.max_budget = 10.0 + existing_key.key_alias = None + existing_key.models = [] + existing_key.metadata = metadata + existing_key.model_dump.return_value = { + "token": token, + "user_id": "internal_user", + "team_id": None, + "max_budget": 10.0, + } + return existing_key + + +def _wire_update_key_fn(monkeypatch, existing_key): + mock_prisma_client = AsyncMock() + updated_key = MagicMock() + updated_key.token = existing_key.token + updated_key.key_alias = "my-alias" + + mock_prisma_client.get_data = AsyncMock(return_value=existing_key) + mock_prisma_client.update_data = AsyncMock(return_value=updated_key) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", lambda token: existing_key.token) + + async def _noop(**kwargs): + pass + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + _noop, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias", + _noop, + ) + + +@pytest.mark.asyncio +async def test_update_key_output_token_estimate_lowered_rejected_for_non_admin(monkeypatch): + """End-to-end wiring: a key's owner reaches /key/update without any admin + check because metadata is a non-budget field, so the gate has to fire + inside the update path itself rather than only in a helper.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {_ESTIMATE: 4000})) + + mock_request = MagicMock() + mock_request.query_params = {} + + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, default_estimated_output_tokens=1), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can set" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_key_output_token_estimate_unchanged_allows_non_admin_edit(monkeypatch): + """The edit form resends every field it renders, so gating on presence + would 403 a key owner renaming their own key.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "b1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {_ESTIMATE: 4000})) + + mock_request = MagicMock() + mock_request.query_params = {} + + result = await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, key_alias="my-alias", default_estimated_output_tokens=4000), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert result is not None + + +@pytest.mark.asyncio +async def test_regenerate_key_output_token_estimate_lowered_rejected_for_non_admin(): + """/key/regenerate is a third write path into the same stored metadata. + + can_modify_verification_token lets a key's own holder regenerate it, and + the request body runs through prepare_key_update_data exactly as an update + does, so gating only generate and update leaves the declaration writable. + """ + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + token = "c1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + key_in_db = LiteLLM_VerificationToken( + token=token, + user_id="internal_user", + metadata={_ESTIMATE: 4000}, + ) + + with pytest.raises(HTTPException) as exc: + await _execute_virtual_key_regeneration( + prisma_client=AsyncMock(), + key_in_db=key_in_db, + hashed_api_key=token, + key="sk-original", + data=RegenerateKeyRequest(key="sk-original", default_estimated_output_tokens=1), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert exc.value.status_code == 403 + assert "Only proxy admins can set" in str(exc.value.detail) diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py new file mode 100644 index 00000000000..30fe78d93c7 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -0,0 +1,709 @@ +"""Tests for PTU config on the model deployment (v1 model-settings design).""" + +import datetime +import json +from contextlib import ExitStack +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LiteLLM_ProxyModelTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.management_endpoints.model_management_endpoints import ( + _merged_ptu_model_info, + _raise_if_ptu_cost_attribution_disabled, + _validate_ptu_model_info, + add_new_model, + update_db_model, +) +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment + + +def test_model_info_accepts_valid_ptu_fields(): + info = ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ) + assert info.ptu_count == 5 + assert info.cost_per_ptu_per_hour == 2.0 + + +def test_model_info_rejects_non_positive_count(): + with pytest.raises(ValueError): + ModelInfo( + id="x", + team_id="t", + ptu_count=0, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ) + + +def test_model_info_rejects_negative_rate(): + with pytest.raises(ValueError): + ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=-1.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ) + + +def test_model_info_rejects_a_count_beyond_the_cap(): + """flat cost multiplies the count by a float, and an unbounded int overflows that + conversion, which aborted the rollup for every team rather than skipping one model.""" + with pytest.raises(ValueError): + ModelInfo(id="x", team_id="t", ptu_count=10**400, cost_per_ptu_per_hour=2.0) + + +def test_model_info_accepts_a_count_at_the_cap(): + info = ModelInfo(id="x", team_id="t", ptu_count=ModelInfo.MAX_PTU_COUNT, cost_per_ptu_per_hour=2.0) + assert info.ptu_count == ModelInfo.MAX_PTU_COUNT + + +@pytest.mark.parametrize("rate", [float("nan"), float("inf"), float("-inf")]) +def test_model_info_rejects_a_non_finite_rate(rate): + """NaN compares False against every bound, so a bare `< 0` check let it through and the + deployment then accrued a flat cost of nan.""" + with pytest.raises(ValueError): + ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=rate) + + +def test_model_info_rejects_a_rate_beyond_the_cap(): + with pytest.raises(ValueError): + ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=ModelInfo.MAX_COST_PER_PTU_PER_HOUR * 2) + + +def test_model_info_allows_partial_delta_for_patch(): + # A PATCH delta may carry only one field; bounds-only validation must not reject it. + info = ModelInfo(id="x", ptu_count=5) + assert info.ptu_count == 5 + assert info.cost_per_ptu_per_hour is None + + +def test_validate_helper_no_ptu_is_noop(): + _validate_ptu_model_info({"team_id": "t"}) + + +def test_validate_helper_requires_both_fields(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info({"team_id": "t", "ptu_count": 5}) + assert exc.value.status_code == 400 + assert "set together" in exc.value.detail + + +def test_validate_helper_requires_team_id(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "ptu_effective_from": "2026-08-01T00:00:00Z"} + ) + assert exc.value.status_code == 400 + assert "team_id" in exc.value.detail + + +def test_validate_helper_requires_an_effective_start(): + """Flat cost accrues from the start, so it cannot be inferred.""" + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info({"team_id": "t", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0}) + assert exc.value.status_code == 400 + assert "ptu_effective_from is required" in exc.value.detail + + +def test_validate_helper_passes_full_config(): + _validate_ptu_model_info( + {"team_id": "t", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "ptu_effective_from": "2026-08-01T00:00:00Z"} + ) + + +def test_model_info_rejects_effective_to_before_from(): + import datetime + + with pytest.raises(ValueError): + ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 30, tzinfo=datetime.timezone.utc), + ptu_effective_to=datetime.datetime(2026, 7, 29, tzinfo=datetime.timezone.utc), + ) + + +def test_model_info_accepts_valid_effective_window(): + import datetime + + info = ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 30, tzinfo=datetime.timezone.utc), + ptu_effective_to=datetime.datetime(2026, 8, 30, tzinfo=datetime.timezone.utc), + ) + assert info.ptu_effective_from is not None + + +def test_model_info_compares_mixed_naive_and_aware_timestamps(): + import datetime + + info = ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 30, 23, 0), + ptu_effective_to=datetime.datetime(2026, 7, 31, 0, 0, tzinfo=datetime.timezone.utc), + ) + assert info.ptu_effective_to is not None + + with pytest.raises(ValueError): + ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 31, 2, 0), + ptu_effective_to=datetime.datetime(2026, 7, 31, 0, 0, tzinfo=datetime.timezone.utc), + ) + + +def test_validate_helper_rejects_effective_to_before_from(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + { + "team_id": "t", + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": "2026-07-30T00:00:00Z", + "ptu_effective_to": "2026-07-29T00:00:00Z", + } + ) + assert exc.value.status_code == 400 + assert "ptu_effective_to" in exc.value.detail + + +def test_validate_helper_accepts_valid_window_on_merged_info(): + _validate_ptu_model_info( + { + "team_id": "t", + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": "2026-07-30T00:00:00Z", + "ptu_effective_to": "2026-08-30T00:00:00Z", + } + ) + + +def test_validate_helper_rejects_inverted_window_without_count_or_rate(): + """A patch that touches only one end of the window merges to a model_info with no count + or rate. Returning early on that shape let an inverted window reach the row, and the next + load then failed to parse it and dropped the deployment out of the router.""" + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + { + "team_id": "t", + "ptu_effective_from": "2026-08-02T00:00:00Z", + "ptu_effective_to": "2026-08-01T00:00:00Z", + } + ) + assert exc.value.status_code == 400 + assert "ptu_effective_to" in exc.value.detail + + +def test_validate_helper_rejects_equal_window_bounds_without_count_or_rate(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + { + "ptu_effective_from": "2026-08-01T00:00:00Z", + "ptu_effective_to": "2026-08-01T00:00:00Z", + } + ) + assert exc.value.status_code == 400 + + +def test_validate_helper_accepts_ordered_window_without_count_or_rate(): + """Window-only edits stay legal; only the ordering is enforced, and no team_id is + demanded while the deployment carries no priced PTU config.""" + _validate_ptu_model_info( + { + "ptu_effective_from": "2026-08-01T00:00:00Z", + "ptu_effective_to": "2026-08-02T00:00:00Z", + } + ) + + +def test_validate_helper_accepts_a_single_open_ended_bound(): + _validate_ptu_model_info({"ptu_effective_from": "2026-08-01T00:00:00Z"}) + _validate_ptu_model_info({"ptu_effective_to": "2026-08-02T00:00:00Z"}) + + +class TestPartialPtuEditsUseTheMergedView: + """A PTU invariant holds over the deployment as it will exist, not over whichever + subset of fields a caller sent. Validating the patch alone rejected an ordinary edit.""" + + @pytest.fixture(autouse=True) + def _enabled(self, monkeypatch): + """PTU writes are gated off by default; these are about the validator, not the gate.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + @staticmethod + def _configured(): + return Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_count=10, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 7, 1, tzinfo=datetime.timezone.utc), + ), + ) + + def test_raising_the_rate_on_a_configured_model_is_allowed(self): + """The patch carries no start; the stored row supplies it.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=10, cost_per_ptu_per_hour=3.0)), + ) + _validate_ptu_model_info(merged) + assert merged["cost_per_ptu_per_hour"] == 3.0 + assert merged["ptu_effective_from"] is not None + + def test_a_genuinely_startless_configuration_is_still_rejected(self): + """Merging must not become a way to smuggle PTU config in without a start.""" + bare = Deployment(model_name="gpt-4o", litellm_params=LiteLLM_Params(model="openai/gpt-4o")) + merged = _merged_ptu_model_info( + db_model=bare, + patch_data=updateDeployment( + model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=10, cost_per_ptu_per_hour=2.0) + ), + ) + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info(merged) + assert "ptu_effective_from is required" in exc.value.detail + + def test_the_patch_still_wins_over_the_stored_value(self): + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=25)), + ) + assert merged["ptu_count"] == 25 + + def test_an_explicit_null_clears_the_stored_field(self): + """update_db_model drops a PTU field a patch sends as null, so the merged view has to + drop it too. Carrying the stored value forward validated a deployment that never + existed.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None)), + ) + assert "ptu_count" not in merged + + def test_clearing_one_half_of_the_pair_is_rejected(self): + """The write leaves a rate with no count. Merging on the stored count hid that.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None)), + ) + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info(merged) + assert "must be set together" in exc.value.detail + + def test_clearing_the_whole_pair_is_allowed(self): + """Turning PTU off on a deployment is a legitimate edit.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None)), + ) + _validate_ptu_model_info(merged) + assert "ptu_count" not in merged + assert "cost_per_ptu_per_hour" not in merged + + def test_an_omitted_field_is_not_a_clear(self): + """A partial edit that never mentions the count keeps it. Only an explicit null clears.""" + merged = _merged_ptu_model_info( + db_model=self._configured(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", cost_per_ptu_per_hour=3.0)), + ) + assert merged["ptu_count"] == 10 + + +class TestTeamModelUpdateValidatesBeforeWriting: + """Drives the endpoint path itself, not the helpers. The validator sits above the team + ACL write, which autocommits, so what it validates has to be right at that call site.""" + + @pytest.fixture(autouse=True) + def _enabled(self, monkeypatch): + """PTU writes are gated off by default; these are about the validator, not the gate.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + @staticmethod + async def _run(db_model, patch_data, monkeypatch, touched=None): + import litellm.proxy.management_endpoints.model_management_endpoints as mme + + touched = [] if touched is None else touched + + async def _never(*args, **kwargs): + touched.append("team_write") + + monkeypatch.setattr(mme, "_setup_new_team_model_assignment", _never) + monkeypatch.setattr(mme, "_update_existing_team_model_assignment", _never) + monkeypatch.setattr(mme.ModelManagementAuthChecks, "allow_team_model_action", AsyncMock(return_value=True)) + result = await mme._update_team_model_in_db( + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=MagicMock(), + prisma_client=MagicMock(), + ) + return result, touched + + @pytest.mark.asyncio + async def test_raising_the_rate_on_a_configured_model_reaches_the_write(self, monkeypatch): + """The patch carries no start. Validating it alone rejected this ordinary edit.""" + db_model = TestPartialPtuEditsUseTheMergedView._configured() + patch = updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=10, cost_per_ptu_per_hour=3.0)) + + result, touched = await self._run(db_model, patch, monkeypatch) + + assert touched == ["team_write"] + assert json.loads(result["model_info"])["cost_per_ptu_per_hour"] == 3.0 + + @pytest.mark.asyncio + async def test_a_startless_configuration_is_refused_before_the_team_write(self, monkeypatch): + """And the refusal still lands before anything is committed.""" + bare = Deployment(model_name="gpt-4o", litellm_params=LiteLLM_Params(model="openai/gpt-4o")) + patch = updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=10, cost_per_ptu_per_hour=2.0)) + + with pytest.raises(HTTPException) as exc: + await self._run(bare, patch, monkeypatch) + + assert "ptu_effective_from is required" in exc.value.detail + + @pytest.mark.asyncio + async def test_the_gate_refuses_before_the_team_write(self, monkeypatch): + """The gate lived inside update_db_model, which runs after the team ACL write, so a + rejected edit still moved the model between teams.""" + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + db_model = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="dep-0", team_id="team-A"), + ) + patch = updateDeployment( + model_info=ModelInfo( + id="dep-0", + team_id="team-B", + ptu_count=15, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 8, 1, tzinfo=datetime.timezone.utc), + ) + ) + touched = [] + + with pytest.raises(HTTPException) as exc: + await self._run(db_model, patch, monkeypatch, touched) + + assert PTU_COST_ATTRIBUTION_ENV_VAR in exc.value.detail + assert touched == [] + + @pytest.mark.asyncio + async def test_clearing_half_the_pair_is_refused_before_the_team_write(self, monkeypatch): + """The write drops the nulled field, so validating against the stored one let a + deployment with a rate and no count commit.""" + db_model = TestPartialPtuEditsUseTheMergedView._configured() + patch = updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=None)) + touched = [] + + with pytest.raises(HTTPException) as exc: + await self._run(db_model, patch, monkeypatch, touched) + + assert "must be set together" in exc.value.detail + assert touched == [] + + @pytest.mark.asyncio + async def test_clearing_the_whole_pair_reaches_the_write_and_stores_neither_field(self, monkeypatch): + """What the validator approved is what the write persists.""" + db_model = TestPartialPtuEditsUseTheMergedView._configured() + patch = updateDeployment( + model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=None, cost_per_ptu_per_hour=None) + ) + + result, touched = await self._run(db_model, patch, monkeypatch) + + assert touched == ["team_write"] + stored = json.loads(result["model_info"]) + assert "ptu_count" not in stored + assert "cost_per_ptu_per_hour" not in stored + + +class TestPtuCostAttributionGate: + """PTU config is only writable once an operator sets LITELLM_ENABLE_PTU_COST_ATTRIBUTION. + + The fields are rejected rather than dropped: a silent accept-and-drop would let a + caller believe a flat cost was configured while the rollup that prices it is not + even scheduled. + """ + + @pytest.fixture(autouse=True) + def _flag_off(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + @pytest.fixture + def flag_on(self, monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + @pytest.mark.parametrize( + "model_info", + [ + {"team_id": "t", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0}, + {"ptu_count": 5}, + {"cost_per_ptu_per_hour": 2.0}, + {"ptu_effective_from": "2026-08-01T00:00:00Z"}, + {"ptu_effective_to": "2026-08-02T00:00:00Z"}, + ], + ) + def test_rejects_any_ptu_field_while_disabled(self, model_info): + with pytest.raises(HTTPException) as exc: + _raise_if_ptu_cost_attribution_disabled(model_info) + assert exc.value.status_code == 400 + assert PTU_COST_ATTRIBUTION_ENV_VAR in exc.value.detail + + def test_names_every_offending_field(self): + with pytest.raises(HTTPException) as exc: + _raise_if_ptu_cost_attribution_disabled({"ptu_count": 5, "cost_per_ptu_per_hour": 2.0}) + assert "ptu_count" in exc.value.detail + assert "cost_per_ptu_per_hour" in exc.value.detail + + def test_allows_a_request_without_ptu_fields_while_disabled(self): + _raise_if_ptu_cost_attribution_disabled({"team_id": "t", "access_groups": ["a"]}) + + def test_allows_every_ptu_field_once_enabled(self, flag_on): + _raise_if_ptu_cost_attribution_disabled( + { + "team_id": "t", + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": "2026-08-01T00:00:00Z", + "ptu_effective_to": "2026-08-02T00:00:00Z", + } + ) + + +def _deployment_without_ptu() -> Deployment: + return Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="dep-0", team_id="t"), + ) + + +def _deployment_with_stored_ptu() -> Deployment: + return Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_count=15, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ), + ) + + +class TestUpdateDbModelPtuGate: + @pytest.fixture(autouse=True) + def _flag_off(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + def test_patch_carrying_ptu_config_is_rejected(self): + with pytest.raises(HTTPException) as exc: + update_db_model( + db_model=_deployment_without_ptu(), + updated_patch=updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=15)), + ) + assert exc.value.status_code == 400 + + def test_patch_that_touches_nothing_ptu_still_succeeds(self): + result = update_db_model( + db_model=_deployment_without_ptu(), + updated_patch=updateDeployment(model_info=ModelInfo(id="dep-0", access_groups=["a"])), + ) + assert json.loads(result["model_info"])["access_groups"] == ["a"] + + def test_unrelated_patch_of_a_model_that_stores_ptu_config_is_not_blocked(self): + """A deployment configured during an earlier opt-in stays editable: the gate reads the + incoming patch, not the merged deployment, so the stored config is left in place.""" + result = update_db_model( + db_model=_deployment_with_stored_ptu(), + updated_patch=updateDeployment(model_name="gpt-4o-renamed"), + ) + assert result["model_name"] == "gpt-4o-renamed" + + def test_explicit_nulls_do_not_erase_stored_ptu_config_while_disabled(self): + """A client round-tripping a model_info blob sends the PTU keys as nulls. While the + feature is disabled those nulls must not reach the clear loop: disabling pauses PTU, + it does not silently discard a billing configuration the operator set up earlier.""" + result = update_db_model( + db_model=_deployment_with_stored_ptu(), + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None) + ), + ) + stored = json.loads(result["model_info"]) + assert stored["ptu_count"] == 15 + assert stored["cost_per_ptu_per_hour"] == 2.0 + + def test_the_merged_view_agrees_with_the_write_while_disabled(self): + """The validator sees what the write will store. If the merged view honoured a null the + clear loop ignores, a round-tripped blob would 400 on a half-set pair that never forms.""" + merged = _merged_ptu_model_info( + db_model=_deployment_with_stored_ptu(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None)), + ) + assert merged["ptu_count"] == 15 + _validate_ptu_model_info(merged) + + def test_explicit_nulls_still_clear_once_enabled(self, monkeypatch): + """Clearing remains available to an operator who opted in, which is how PTU config is + removed from a deployment.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + result = update_db_model( + db_model=_deployment_with_stored_ptu(), + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None) + ), + ) + stored = json.loads(result["model_info"]) + assert "ptu_count" not in stored + assert "cost_per_ptu_per_hour" not in stored + + def test_patch_carrying_ptu_config_is_accepted_once_enabled(self, monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + result = update_db_model( + db_model=_deployment_without_ptu(), + updated_patch=updateDeployment( + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_count=15, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ) + ), + ) + stored = json.loads(result["model_info"]) + assert stored["ptu_count"] == 15 + assert stored["cost_per_ptu_per_hour"] == 2.0 + + +class TestAddNewModelPtuGate: + @pytest.fixture(autouse=True) + def _flag_off(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + @staticmethod + def _patched_proxy(model_id: str): + """Patch everything /model/new touches except the PTU gate, and hand back the DB writers.""" + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="ptu-model", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={"id": model_id}, + created_by="test-admin", + updated_by="test-admin", + ) + add_model_to_db = AsyncMock(return_value=db_row) + add_team_model_to_db = AsyncMock(return_value=db_row) + + mock_proxy_config = MagicMock() + mock_proxy_config.add_deployment = AsyncMock(return_value=None) + + mock_router = MagicMock() + mock_router.get_model_ids.return_value = [model_id] + + proxy_server = "litellm.proxy.proxy_server" + endpoints = "litellm.proxy.management_endpoints.model_management_endpoints" + return (add_model_to_db, add_team_model_to_db), [ + patch(f"{proxy_server}.prisma_client", MagicMock()), + patch(f"{proxy_server}.store_model_in_db", True), + patch(f"{proxy_server}.proxy_config", mock_proxy_config), + patch(f"{proxy_server}.proxy_logging_obj", MagicMock()), + patch(f"{proxy_server}.general_settings", {}), + patch(f"{proxy_server}.premium_user", True), + patch(f"{proxy_server}.llm_router", mock_router), + patch( + f"{endpoints}.ModelManagementAuthChecks.can_user_make_model_call", + AsyncMock(return_value=True), + ), + patch(f"{endpoints}._add_model_to_db", add_model_to_db), + patch(f"{endpoints}._add_team_model_to_db", add_team_model_to_db), + ] + + @staticmethod + def _ptu_deployment(model_id: str) -> Deployment: + return Deployment( + model_name="ptu-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"), + model_info=ModelInfo( + id=model_id, + team_id="team-1", + ptu_count=15, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ), + ) + + @pytest.mark.asyncio + async def test_model_new_rejects_ptu_config_while_disabled(self): + (add_model_to_db, add_team_model_to_db), patches = self._patched_proxy("ptu-gate-model") + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + with pytest.raises(Exception) as exc: + await add_new_model(model_params=self._ptu_deployment("ptu-gate-model"), user_api_key_dict=admin) + + assert PTU_COST_ATTRIBUTION_ENV_VAR in str(exc.value) + add_model_to_db.assert_not_called() + add_team_model_to_db.assert_not_called() + + @pytest.mark.asyncio + async def test_model_new_accepts_a_deployment_without_ptu_config_while_disabled(self): + _, patches = self._patched_proxy("plain-model") + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + result = await add_new_model( + model_params=Deployment( + model_name="ptu-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"), + model_info=ModelInfo(id="plain-model"), + ), + user_api_key_dict=admin, + ) + + assert result.model_id == "plain-model" + + @pytest.mark.asyncio + async def test_model_new_accepts_ptu_config_once_enabled(self, monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + (_, add_team_model_to_db), patches = self._patched_proxy("ptu-gate-model") + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + result = await add_new_model(model_params=self._ptu_deployment("ptu-gate-model"), user_api_key_dict=admin) + + assert result.model_id == "ptu-gate-model" + add_team_model_to_db.assert_called_once() diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 1e47010b57c..6abc40eb28e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -380,6 +380,66 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth): app.dependency_overrides = {} +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["budget_duration", "team_member_budget_duration"]) +@pytest.mark.parametrize("bad_duration", ["0s", "-5m"]) +async def test_new_team_rejects_a_duration_that_never_advances( + mock_db_client, mock_admin_auth, field, bad_duration +): + """A zero-length window resets to "now", so the team row is due again the + moment it is written. The reset job re-reads such rows on every tick, and a + tenant with enough of them fills each batch and starves other tenants. + """ + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + mock_db_client.db = MagicMock() + mock_team_create = AsyncMock() + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = mock_team_create + + with pytest.raises(ProxyException) as exc_info: + await new_team( + data=NewTeamRequest(team_alias="my-team", **{field: bad_duration}), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert str(exc_info.value.code) == "400" + assert "Invalid budget_duration" in str(exc_info.value.message) + mock_team_create.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ["budget_duration", "team_member_budget_duration"]) +async def test_update_team_rejects_a_duration_that_never_advances( + mock_db_client, mock_admin_auth, field +): + """/team/update must reject the same never-advancing durations /team/new does.""" + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import update_team + + mock_db_client.db = MagicMock() + mock_find_unique = AsyncMock(return_value=None) + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.find_unique = mock_find_unique + + with pytest.raises(ProxyException) as exc_info: + await update_team( + data=UpdateTeamRequest(team_id="team-1", **{field: "0s"}), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert str(exc_info.value.code) == "400" + assert "Invalid budget_duration" in str(exc_info.value.message) + mock_find_unique.assert_not_awaited() + + @pytest.mark.asyncio async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): """ @@ -11008,3 +11068,207 @@ def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back(): assert f"u{_MAX_REPORTED_UNKNOWN_USER_IDS}" not in detail assert f"and {500 - _MAX_REPORTED_UNKNOWN_USER_IDS} more" in detail assert len(detail) < 1000 + + +_TEAM_ESTIMATE = "default_estimated_output_tokens" +_TEAM_ESTIMATE_PER_MODEL = "default_estimated_output_tokens_per_model" + + +@pytest.mark.parametrize( + "label, request_body, existing_metadata, allowed", + [ + ("nothing declared", {}, None, True), + ("declared top-level with none stored", {_TEAM_ESTIMATE: 1}, None, False), + ("declared inside metadata with none stored", {"metadata": {_TEAM_ESTIMATE: 1}}, None, False), + ( + "per-model map declared inside metadata", + {"metadata": {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 1}}}, + None, + False, + ), + ("unrelated edit, metadata omitted", {"tpm_limit": 99}, {_TEAM_ESTIMATE: 2000}, True), + ("stored value resent unchanged", {_TEAM_ESTIMATE: 2000}, {_TEAM_ESTIMATE: 2000}, True), + ("stored value lowered", {_TEAM_ESTIMATE: 1}, {_TEAM_ESTIMATE: 2000}, False), + ("stored value raised", {_TEAM_ESTIMATE: 9000}, {_TEAM_ESTIMATE: 2000}, False), + ( + "stored value cleared by sending a metadata blob without it", + {"metadata": {"other": "keep"}}, + {_TEAM_ESTIMATE: 2000, "other": "keep"}, + False, + ), + ( + "stored value resent inside the metadata blob", + {"metadata": {_TEAM_ESTIMATE: 2000, "other": "keep"}}, + {_TEAM_ESTIMATE: 2000, "other": "keep"}, + True, + ), + ( + "per-model map resent unchanged", + {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + True, + ), + ( + "one model in the per-model map lowered", + {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 1}}, + {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + False, + ), + ], +) +def test_team_output_token_estimate_admin_gate_matrix(label, request_body, existing_metadata, allowed): + """A team admin may only leave a team's stored output-token estimate exactly as it is. + + A team admin can write team metadata, and every key on the team inherits the + team declaration, so without this a team admin could shrink the reservation + for the whole team and under-reserve against an organization TPM window the + organization set above them. Same value-transition rule as the key gate, + including the raw-metadata route and clearing by omission. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.auth.auth_utils import ( + enforce_output_token_estimates_are_admin_only, + ) + + def _call(caller): + enforce_output_token_estimates_are_admin_only( + data=UpdateTeamRequest(team_id="t", **request_body), + existing_metadata=existing_metadata, + user_api_key_dict=caller, + entity="team", + ) + + team_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-team-admin", + user_id="team-admin", + ) + if allowed: + _call(team_admin) + else: + with pytest.raises(HTTPException) as exc: + _call(team_admin) + assert exc.value.status_code == 403 + assert "on a team" in str(exc.value.detail) + + _call( + UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin", + ) + ) + + +def _wire_update_team(stack, existing_metadata): + """Mock just enough of update_team to reach (or pass) the estimate gate.""" + from unittest.mock import AsyncMock, MagicMock, patch + + mock_prisma_client = stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client")) + stack.enter_context(patch("litellm.proxy.proxy_server.llm_router")) + stack.enter_context(patch("litellm.proxy.proxy_server.user_api_key_cache")) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) + stack.enter_context(patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")) + stack.enter_context(patch("litellm.proxy.management_endpoints.team_endpoints._cache_team_object")) + + existing_team = MagicMock() + existing_team.metadata = existing_metadata + existing_team.model_dump.return_value = { + "team_id": "test_team_id", + "team_alias": "test_team", + "metadata": existing_metadata, + "members_with_roles": [{"user_id": "team-admin", "role": "admin"}], + } + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + + updated_team = MagicMock() + updated_team.team_id = "test_team_id" + updated_team.model_dump.return_value = {"team_id": "test_team_id"} + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) + mock_prisma_client.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_update_team_output_token_estimate_lowered_rejected_for_team_admin(): + """End-to-end wiring: _verify_team_access admits a team admin, so the gate + has to fire inside update_team itself.""" + import contextlib + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with contextlib.ExitStack() as stack: + _wire_update_team(stack, {_TEAM_ESTIMATE: 4000}) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", default_estimated_output_tokens=1), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-team-admin", + user_id="team-admin", + ), + ) + + assert str(exc.value.code) == "403" + assert "on a team" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_team_output_token_estimate_unchanged_allows_team_admin_edit(): + """The team settings form resends every field it renders, so gating on + presence would break a team admin editing an unrelated setting.""" + import contextlib + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {_TEAM_ESTIMATE: 4000}) + await update_team( + data=UpdateTeamRequest( + team_id="test_team_id", + team_alias="renamed", + default_estimated_output_tokens=4000, + ), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-team-admin", + user_id="team-admin", + ), + ) + + assert prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_new_team_output_token_estimate_rejected_for_non_admin(): + """/team/new is the other write path into the same stored declaration.""" + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with pytest.raises(ProxyException) as exc: + await new_team( + data=NewTeamRequest(team_alias="t", default_estimated_output_tokens=1), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + ) + + assert str(exc.value.code) == "403" + assert "on a team" in str(exc.value.message) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py b/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py new file mode 100644 index 00000000000..07a85a70815 --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_storage_backend_service.py @@ -0,0 +1,127 @@ +import pytest + +from litellm.llms.base_llm.files.transformation import BaseFileEndpoints +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.openai_files_endpoints import storage_backend_service +from litellm.proxy.openai_files_endpoints.storage_backend_service import ( + StorageBackendFileService, +) + + +class _RecordingStorageBackend: + def __init__(self): + self.upload_calls = [] + + async def upload_file(self, **kwargs): + self.upload_calls.append(kwargs) + return "https://storage.example/blob-1" + + +class _FakeManagedFilesHook(BaseFileEndpoints): + def __init__(self): + self.stored = [] + + async def acreate_file( + self, create_file_request, llm_router, target_model_names_list, litellm_parent_otel_span, user_api_key_dict + ): + raise NotImplementedError + + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router=None): + raise NotImplementedError + + async def afile_list(self, purpose, litellm_parent_otel_span, **data): + raise NotImplementedError + + async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + raise NotImplementedError + + async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + raise NotImplementedError + + async def store_unified_file_id(self, **kwargs): + self.stored.append(kwargs) + + +class _FakeProxyLogging: + def __init__(self, hook): + self._hook = hook + + def get_proxy_hook(self, hook_name): + return self._hook if hook_name == "managed_files" else None + + +def _file_data(): + return {"content": b"x", "filename": "input.jsonl", "content_type": "application/jsonl"} + + +@pytest.mark.asyncio +async def test_upload_with_target_model_names_but_no_hook_raises_before_uploading(monkeypatch): + backend = _RecordingStorageBackend() + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend) + + with pytest.raises(ProxyException) as exc_info: + await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=["gpt-x"], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=None), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + snapshot = { + "code": exc_info.value.code, + "message_names_requirement": "requires a database-connected proxy" in exc_info.value.message, + "upload_calls": backend.upload_calls, + } + assert snapshot == {"code": "400", "message_names_requirement": True, "upload_calls": []} + + +@pytest.mark.asyncio +async def test_upload_without_target_model_names_skips_hook_requirement(monkeypatch): + backend = _RecordingStorageBackend() + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend) + + file_object = await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=[], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=None), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + snapshot = { + "upload_count": len(backend.upload_calls), + "id_prefix": file_object.id.split("-")[0], + } + assert snapshot == {"upload_count": 1, "id_prefix": "file"} + + +@pytest.mark.asyncio +async def test_upload_with_target_model_names_and_hook_stores_unified_id(monkeypatch): + backend = _RecordingStorageBackend() + monkeypatch.setattr(storage_backend_service, "get_storage_backend", lambda name: backend) + hook = _FakeManagedFilesHook() + + file_object = await StorageBackendFileService.upload_file_to_storage_backend( + file_data=_file_data(), + target_storage="azure_storage", + target_model_names=["gpt-x"], + purpose="batch", + proxy_logging_obj=_FakeProxyLogging(hook=hook), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + snapshot = { + "upload_count": len(backend.upload_calls), + "store_count": len(hook.stored), + "stored_id_matches_response": hook.stored[0]["file_id"] == file_object.id, + "model_mappings": hook.stored[0]["model_mappings"], + } + assert snapshot == { + "upload_count": 1, + "store_count": 1, + "stored_id_matches_response": True, + "model_mappings": {"gpt-x": "https://storage.example/blob-1"}, + } 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 181846fe289..f631215c03d 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 @@ -2,6 +2,7 @@ import json import os import sys import traceback +from typing import Final from unittest import mock from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -2814,6 +2815,63 @@ class TestOpenAIPassthroughRoute: assert result == {"id": "asst_123", "object": "assistant"} +def _resolve_route_name(method: str, path: str) -> str | None: + from starlette.routing import Match + + from litellm.proxy.proxy_server import app + + scope: Final = { + "type": "http", + "method": method, + "path": path, + "headers": [], + "query_string": b"", + "root_path": "", + } + for route in app.router.routes: + if route.matches(scope)[0] == Match.FULL: + return getattr(route, "name", None) + return None + + +@pytest.mark.parametrize( + "method, path", + [ + ("POST", "/openai_passthrough/v1/files"), + ("GET", "/openai_passthrough/v1/files"), + ("GET", "/openai_passthrough/v1/files/file-abc123"), + ("DELETE", "/openai_passthrough/v1/files/file-abc123"), + ("GET", "/openai_passthrough/v1/files/file-abc123/content"), + ("POST", "/openai_passthrough/v1/batches"), + ("GET", "/openai_passthrough/v1/batches"), + ("GET", "/openai_passthrough/v1/batches/batch_abc123"), + ("POST", "/openai_passthrough/v1/batches/batch_abc123/cancel"), + ("POST", "/openai_passthrough/v1/responses"), + ], +) +def test_openai_passthrough_prefix_wins_over_native_provider_routes(method, path): + """ + /openai_passthrough exists to guarantee passthrough, so the native + /{provider}/v1/files and /{provider}/v1/batches routes must never capture it + with provider="openai_passthrough" (which 500s on the LlmProviders lookup). + """ + assert _resolve_route_name(method, path) == "openai_proxy_route" + + +@pytest.mark.parametrize( + "method, path, expected_name", + [ + ("POST", "/openai/v1/files", "create_file"), + ("GET", "/azure/v1/files", "list_files"), + ("POST", "/v1/files", "create_file"), + ("POST", "/v1/batches", "create_batch"), + ("POST", "/openai/v1/chat/completions", "openai_proxy_route"), + ], +) +def test_native_provider_routes_are_unchanged(method, path, expected_name): + assert _resolve_route_name(method, path) == expected_name + + class TestCursorProxyRoute: """Tests for the Cursor Cloud Agents pass-through route.""" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index 52da7a4a81d..d53e6dedf0b 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -258,6 +258,7 @@ class TestVertexAIBatchPassthroughHandler: batch_object=batch_object, model_object_id=model_object_id, logging_obj=mock_logging_obj, + is_batch_create=True, user_api_key_dict={"user_id": "test-user"}, ) @@ -307,6 +308,7 @@ class TestVertexAIBatchPassthroughHandler: batch_object={"id": "b1", "object": "batch", "status": "validating"}, model_object_id="b1", logging_obj=mock_logging_obj, + is_batch_create=True, **kwargs, ) @@ -315,6 +317,140 @@ class TestVertexAIBatchPassthroughHandler: assert call_kwargs["user_api_key_dict"].user_id == expected_user_id assert call_kwargs["user_api_key_dict"].team_id == expected_team_id + def _store_with_metadata(self, mock_logging_obj, mock_managed_files_hook, metadata): + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl, + patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.verbose_proxy_logger" + ), + ): + mock_pl.get_proxy_hook.return_value = mock_managed_files_hook + VertexPassthroughLoggingHandler._store_batch_managed_object( + unified_object_id="uoi", + batch_object={"id": "b1", "object": "batch", "status": "validating"}, + model_object_id="b1", + logging_obj=mock_logging_obj, + is_batch_create=True, + litellm_params={"metadata": metadata}, + ) + mock_managed_files_hook.store_unified_object_id.assert_called_once() + return mock_managed_files_hook.store_unified_object_id.call_args[1] + + def test_create_persists_key_hash_and_tags( + self, mock_logging_obj, mock_managed_files_hook + ): + """Regression (spend loss): the batch create must persist the creating key's hashed + token and its tags so CheckBatchCost can attribute the batch-cost spend row. Before + this fix the stored api_key was always "" and the row was dropped as unattributed.""" + call_kwargs = self._store_with_metadata( + mock_logging_obj, + mock_managed_files_hook, + { + "user_api_key": "hashed-key-a", + "user_api_key_user_id": "alice", + "user_api_key_team_id": "team-alpha", + "user_api_key_auth_metadata": {"tags": ["env:prod", 7, "team:ml"]}, + }, + ) + + assert call_kwargs["user_api_key_dict"].api_key == "hashed-key-a" + # non-string tags are dropped so downstream tag budgets cannot be bypassed + assert call_kwargs["request_tags"] == ("env:prod", "team:ml") + assert call_kwargs["persist_attribution"] is True + + @pytest.mark.parametrize( + "metadata, expected", + [ + # a request that sent its own tags (x-litellm-tags header or body metadata) + ({"tags": ["req:a", "req:b"]}, ("req:a", "req:b")), + # request tags win over the key's own tags + ( + {"tags": ["req:a"], "user_api_key_auth_metadata": {"tags": ["key:b"]}}, + ("req:a",), + ), + # no request tags: fall back to the tags the key itself carries + ({"user_api_key_auth_metadata": {"tags": ["key:b"]}}, ("key:b",)), + # neither: no tags on the spend row + ({}, None), + ], + ) + def test_request_tags_precedence( + self, mock_logging_obj, mock_managed_files_hook, metadata, expected + ): + """Request tags take precedence over the key's tags, and the key's tags are the + fallback because a tagged key does not put its tags in the top-level metadata.""" + call_kwargs = self._store_with_metadata( + mock_logging_obj, + mock_managed_files_hook, + {"user_api_key": "hashed-key-a", **metadata}, + ) + + assert call_kwargs["request_tags"] == expected + + @pytest.mark.parametrize( + "url_route, expected", + [ + ("/v1/projects/p/locations/us-central1/batchPredictionJobs", True), + ("/v1/projects/p/locations/us-central1/batchPredictionJobs/", True), + ("/v1/projects/p/locations/us-central1/batchPredictionJobs?alt=json", True), + ("/v1/projects/p/locations/us-central1/batchPredictionJobs/123456", False), + ("/v1/projects/p/locations/us-central1/batchPredictionJobs/123456?alt=json", False), + ], + ) + def test_batch_is_registered_from_the_create_route_only( + self, mock_logging_obj, url_route, expected + ): + """Only a POST to the collection route is the create, and only the create claims + attribution. Every id-scoped route is a poll or retrieve, which still reports the + batch so its status and file object stay in sync, but carries is_batch_create=False + so it neither claims the batch nor creates a row it would then own.""" + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "name": "projects/p/locations/us-central1/batchPredictionJobs/123456", + "model": "publishers/google/models/gemini-2.5-flash", + } + + with ( + patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.verbose_proxy_logger" + ), + patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.VertexPassthroughLoggingHandler._store_batch_managed_object" + ) as mock_store, + patch( + "litellm.llms.vertex_ai.batches.transformation.VertexAIBatchTransformation" + ) as mock_transformation, + patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.VertexPassthroughLoggingHandler.get_actual_model_id_from_router", + return_value="gemini-2.5-flash", + ), + ): + mock_transformation.transform_vertex_ai_batch_response_to_openai_batch_response.return_value = { + "id": "123456", + "object": "batch", + "status": "validating", + "created_at": 1704067200, + "input_file_id": "gs://bucket/in.jsonl", + "completion_window": "24h", + } + mock_transformation._get_batch_id_from_vertex_ai_batch_response.return_value = "123456" + + VertexPassthroughLoggingHandler.batch_prediction_jobs_handler( + httpx_response=response, + logging_obj=mock_logging_obj, + url_route=url_route, + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + ) + + # every route reports the batch; only the create claims it + mock_store.assert_called_once() + assert mock_store.call_args[1]["unified_object_id"] + assert mock_store.call_args[1]["is_batch_create"] is expected + def test_batch_cost_calculation_integration(self): """Single Vertex AI response → non-zero cost with correct token counts.""" from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage 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 91a7e1bc2c2..f70be17eb95 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2639,3 +2639,67 @@ async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembe assert clean_agent_registry.config_agents == () clean_agent_registry.load_agents_from_db_and_config(db_agents=None) assert clean_agent_registry.get_agent_list() == () + + +# --------------------------------------------------------------------------- +# _init_guardrails_in_db +# --------------------------------------------------------------------------- + + +def _db_guardrail_row(guardrail_id: str, guardrail_type: str) -> dict[str, object]: + return { + "guardrail_id": guardrail_id, + "guardrail_name": f"name-{guardrail_id}", + "litellm_params": {"guardrail": guardrail_type, "mode": "pre_call"}, + "guardrail_info": None, + "team_id": None, + } + + +@pytest.mark.asyncio +async def test_ProxyConfig__init_guardrails_in_db_skips_only_the_unloadable_row(monkeypatch): + """ + A single DB row that fails to initialize used to abort the whole loop, so one + typo'd guardrail type left the proxy running with zero guardrails loaded. + + The failing row's id must still reach reconcile_db_guardrails so that eviction + pass cannot treat a row that is alive in the DB as one that was deleted. + """ + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy.guardrails import guardrail_registry as registry_module + from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams + + class _RecordingHandler(registry_module.InMemoryGuardrailHandler): + def __init__(self) -> None: + super().__init__() + self.reconciled_with: list[set[str]] = [] + + def reconcile_db_guardrails(self, db_guardrail_ids: set[str]) -> list[str]: + self.reconciled_with.append(set(db_guardrail_ids)) + return super().reconcile_db_guardrails(db_guardrail_ids) + + handler = _RecordingHandler() + monkeypatch.setattr(registry_module, "IN_MEMORY_GUARDRAIL_HANDLER", handler) + + def _initializer(litellm_params: LitellmParams, guardrail: Guardrail) -> CustomGuardrail: + return CustomGuardrail( + guardrail_name=guardrail["guardrail_name"], + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ) + + monkeypatch.setitem(registry_module.guardrail_initializer_registry, "lit5367_ok", _initializer) + + prisma_client = MagicMock() + prisma_client.db.litellm_guardrailstable.find_many = AsyncMock( + return_value=[ + _db_guardrail_row("first", "lit5367_ok"), + _db_guardrail_row("broken", "litellm_tool_permission"), + _db_guardrail_row("last", "lit5367_ok"), + ] + ) + + await ProxyConfig()._init_guardrails_in_db(prisma_client=prisma_client) + + assert sorted(handler.IN_MEMORY_GUARDRAILS) == ["first", "last"] + assert handler.reconciled_with == [{"first", "broken", "last"}] diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index f7e2d276a2e..15758c595c0 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -196,9 +196,7 @@ async def test_async_assistants_data_generator_hook_failure_yields_error_chunk( async def _noop_failure(*args, **kwargs): return None - monkeypatch.setattr( - ps.proxy_logging_obj, "async_post_call_streaming_hook", _boom_hook - ) + monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_hook", _boom_hook) monkeypatch.setattr(ps.proxy_logging_obj, "post_call_failure_hook", _noop_failure) stream = _FakeAssistantsStream([_simple_chunk()]) @@ -385,9 +383,7 @@ def test_get_streaming_fallback_metadata_no_additional_headers(): def test_get_streaming_fallback_metadata_zero_fallback_count(): stream = _FakeStream( [], - hidden_params={ - "additional_headers": {"x-litellm-attempted-fallbacks": 0} - }, + hidden_params={"additional_headers": {"x-litellm-attempted-fallbacks": 0}}, ) assert _get_streaming_fallback_metadata(stream) == (False, None, []) @@ -558,9 +554,7 @@ async def test_apply_streaming_chunk_hooks_appends_to_str_so_far(monkeypatch): async def _passthrough(*, user_api_key_dict, response, data, str_so_far=None): return response - monkeypatch.setattr( - ps.proxy_logging_obj, "async_post_call_streaming_hook", _passthrough - ) + monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_hook", _passthrough) new_chunk, new_str = await _apply_streaming_chunk_hooks( chunk=chunk, @@ -870,9 +864,7 @@ async def test_async_data_generator_mid_stream_exception_yields_error_payload( out.append(line) # First entry is the successful "partial" chunk (bytes), last is the error. - assert any( - isinstance(item, str) and item.startswith('data: {"error":') for item in out - ) + assert any(isinstance(item, str) and item.startswith('data: {"error":') for item in out) # --------------------------------------------------------------------------- @@ -914,3 +906,694 @@ def test_select_data_generator_missing_required_kwarg_raises_type_error(): streaming starts.""" with pytest.raises(TypeError): select_data_generator(response=_async_iter([]), user_api_key_dict=_user_auth()) # type: ignore[call-arg] + + +# --------------------------------------------------------------------------- +# SSE keepalive helpers +# --------------------------------------------------------------------------- + + +from litellm.proxy.proxy_server import ( # noqa: E402 + _iter_with_keepalive, + _keepalive_from_deployment_config, + _make_keepalive_resolver, + _resolve_keepalive_seconds, +) +from litellm.proxy.proxy_server import _KEEPALIVE_MAX_SECONDS, _KEEPALIVE_MIN_SECONDS # noqa: E402 + + +@pytest.mark.asyncio +async def test_iter_with_keepalive_hot_path_no_task_wrapping(): + """When keepalive_seconds <= 0, the generator is a transparent pass-through.""" + chunks = [_simple_chunk(content="a"), _simple_chunk(content="b")] + out = [] + async for item in _iter_with_keepalive(_async_iter(chunks), lambda _: 0, keepalive_seconds=0): + out.append(item) + + assert out == chunks + assert ps._STREAM_KEEPALIVE not in out + + +@pytest.mark.asyncio +async def test_iter_with_keepalive_emits_sentinel_when_stream_stalls(): + """With a short keepalive interval and a stalled upstream, _STREAM_KEEPALIVE + sentinels appear before the delayed chunk arrives. The resolver returns a + constant interval, since this test pins the timing mechanics, not + re-resolution.""" + import asyncio + + async def _slow_stream(): + yield _simple_chunk(content="first") + await asyncio.sleep(0.3) + yield _simple_chunk(content="second") + + items = [] + async for item in _iter_with_keepalive(_slow_stream(), lambda _: 0.05, keepalive_seconds=0.05): + items.append(item) + + sentinels = [i for i in items if i is ps._STREAM_KEEPALIVE] + real_chunks = [i for i in items if i is not ps._STREAM_KEEPALIVE] + + assert len(sentinels) >= 2, f"expected >= 2 sentinels during 0.3s stall; got {len(sentinels)}" + assert len(real_chunks) == 2 + assert real_chunks[0].choices[0].delta.content == "first" + assert real_chunks[1].choices[0].delta.content == "second" + + +@pytest.mark.asyncio +async def test_iter_with_keepalive_cancel_on_early_close(): + """Closing the generator early cancels the in-flight task without raising.""" + import asyncio + + async def _infinite_stream(): + while True: + await asyncio.sleep(10) + yield _simple_chunk() + + gen = _iter_with_keepalive(_infinite_stream(), lambda _: 0.05, keepalive_seconds=0.05) + # Advance once to get the sentinel; then close before the real chunk. + first = await gen.__anext__() + assert first is ps._STREAM_KEEPALIVE + # aclose must not raise, and must drain the cancelled task cleanly. + await gen.aclose() + + +@pytest.mark.asyncio +async def test_iter_with_keepalive_disables_after_fallback_lowers_interval(): + """Greptile P1: a mid-stream router fallback can hand off to a deployment + with a different (or disabled) keepalive policy partway through the same + stream. The interval must be re-resolved against each chunk's own identity, + not the one picked before iteration started, or heartbeats keep using the + pre-fallback deployment's policy for the rest of the stream.""" + import asyncio + + async def _slow_stream(): + yield _simple_chunk(content="first") + await asyncio.sleep(0.3) + yield _simple_chunk(content="second") + + def _resolver(item): + # First chunk resolves under the enabled interval used to start the + # wrapper; every chunk after that resolves as if a fallback disabled it. + return 0.0 if item.choices[0].delta.content == "first" else 999.0 + + items = [] + async for item in _iter_with_keepalive(_slow_stream(), _resolver, keepalive_seconds=0.05): + items.append(item) + + sentinels = [i for i in items if i is ps._STREAM_KEEPALIVE] + real_chunks = [i for i in items if i is not ps._STREAM_KEEPALIVE] + + assert sentinels == [], f"expected no sentinels once the resolver disables keepalive; got {len(sentinels)}" + assert len(real_chunks) == 2 + assert real_chunks[0].choices[0].delta.content == "first" + assert real_chunks[1].choices[0].delta.content == "second" + + +@pytest.mark.asyncio +async def test_iter_with_keepalive_enables_after_fallback_raises_interval(): + """Symmetric case: a mid-stream fallback to a deployment with a *shorter* + keepalive interval must take effect immediately, not stay pinned to the + longer interval the stream started with. The interval used to wait for a + chunk is resolved from the *previous* chunk (the only one seen so far when + that wait begins), so the stall has to follow the fallback chunk rather + than precede it: waiting for "third" is where the shorter interval bites.""" + import asyncio + + async def _slow_stream(): + yield _simple_chunk(content="first") + yield _simple_chunk(content="second") + await asyncio.sleep(0.3) + yield _simple_chunk(content="third") + + def _resolver(item): + # "first" resolves under an interval too long to fire before "second" + # arrives; "second" (the fallback chunk) resolves as if the fallback + # deployment enabled a much shorter interval for everything after it. + return 999.0 if item.choices[0].delta.content == "first" else 0.05 + + items = [] + async for item in _iter_with_keepalive(_slow_stream(), _resolver, keepalive_seconds=999.0): + items.append(item) + + sentinels = [i for i in items if i is ps._STREAM_KEEPALIVE] + real_chunks = [i for i in items if i is not ps._STREAM_KEEPALIVE] + + assert len(sentinels) >= 2, ( + f"expected >= 2 sentinels once the resolver enables a short interval; got {len(sentinels)}" + ) + assert len(real_chunks) == 3 + + +@pytest.mark.asyncio +async def test_iter_with_keepalive_activates_from_a_fully_disabled_start(): + """Greptile P1: a stream can start on a deployment with keepalive off + (keepalive_seconds passed in as 0, not merely a long interval) and fall back + mid-stream to one that enables it. The 0-second start must not be treated as + a one-time decision to skip heartbeats for the rest of the stream: no task + is created while inactive, but every chunk still re-resolves so the fallback + chunk can switch the stream into task-wrapped mode.""" + import asyncio + + async def _slow_stream(): + yield _simple_chunk(content="first") + yield _simple_chunk(content="second") + await asyncio.sleep(0.3) + yield _simple_chunk(content="third") + + def _resolver(item): + # "first" resolves to stay off; "second" (the fallback chunk) resolves + # as if the fallback deployment newly enabled a short interval. + return 0.0 if item.choices[0].delta.content == "first" else 0.05 + + items = [] + async for item in _iter_with_keepalive(_slow_stream(), _resolver, keepalive_seconds=0): + items.append(item) + + sentinels = [i for i in items if i is ps._STREAM_KEEPALIVE] + real_chunks = [i for i in items if i is not ps._STREAM_KEEPALIVE] + + assert len(sentinels) >= 2, ( + f"expected >= 2 sentinels once the resolver activates from a disabled start; got {len(sentinels)}" + ) + assert len(real_chunks) == 3 + + +def test_resolve_keepalive_seconds_client_value_ignored_without_override_permission(monkeypatch): + """keepalive_seconds is operator-only by default: a deployment that hasn't set + allow_client_keepalive_override must not let a client's request-level value + change its behavior at all, since that would let any authenticated client + unilaterally enable heartbeats (and the LB-idle-timeout evasion that comes + with them) for a deployment that never opted in.""" + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = 15.0 + deployment.litellm_params.allow_client_keepalive_override = False + + router = MagicMock() + router.get_deployment.return_value = deployment + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "deploy-locked"} + + result = _resolve_keepalive_seconds({"model": "my-model", "keepalive_seconds": 1}, response=response) + assert result == 15.0 + + +def test_resolve_keepalive_seconds_request_value_wins_when_override_allowed(monkeypatch): + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = None + deployment.litellm_params.allow_client_keepalive_override = True + + router = MagicMock() + router.get_deployment.return_value = deployment + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "deploy-opt-in"} + + result = _resolve_keepalive_seconds({"model": "my-model", "keepalive_seconds": 30}, response=response) + assert result == 30.0 + + +def test_resolve_keepalive_seconds_explicit_zero_disables_when_override_allowed(monkeypatch): + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = 20.0 + deployment.litellm_params.allow_client_keepalive_override = True + + router = MagicMock() + router.get_deployment.return_value = deployment + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "deploy-opt-in"} + + result = _resolve_keepalive_seconds({"model": "my-model", "keepalive_seconds": 0}, response=response) + assert result == 0.0 + + +def test_resolve_keepalive_seconds_clamps_below_minimum(monkeypatch): + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = None + deployment.litellm_params.allow_client_keepalive_override = True + + router = MagicMock() + router.get_deployment.return_value = deployment + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "deploy-opt-in"} + + result = _resolve_keepalive_seconds({"model": "my-model", "keepalive_seconds": 0.001}, response=response) + assert result == _KEEPALIVE_MIN_SECONDS + + +def test_resolve_keepalive_seconds_clamps_above_maximum(monkeypatch): + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = None + deployment.litellm_params.allow_client_keepalive_override = True + + router = MagicMock() + router.get_deployment.return_value = deployment + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "deploy-opt-in"} + + result = _resolve_keepalive_seconds({"model": "my-model", "keepalive_seconds": 9999}, response=response) + assert result == _KEEPALIVE_MAX_SECONDS + + +def test_resolve_keepalive_seconds_non_numeric_returns_zero(monkeypatch): + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = None + deployment.litellm_params.allow_client_keepalive_override = True + + router = MagicMock() + router.get_deployment.return_value = deployment + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "deploy-opt-in"} + + result = _resolve_keepalive_seconds({"model": "my-model", "keepalive_seconds": "not-a-number"}, response=response) + assert result == 0.0 + + +def test_resolve_keepalive_seconds_absent_returns_zero(monkeypatch): + monkeypatch.setattr(ps, "llm_router", None) + result = _resolve_keepalive_seconds({}, response=None) + assert result == 0.0 + + +def test_resolve_keepalive_seconds_deployment_disable_cannot_be_overridden_by_request(monkeypatch): + """A deployment that explicitly sets keepalive_seconds: 0 is a hard operator + disable: an authenticated client must not be able to re-enable heartbeats for + that deployment by passing a positive value in the request body, since that + would let a client evade the deployment's idle-timeout behavior at will. This + holds even if the deployment also grants override permission, since an + explicit disable is a stronger, unconditional signal than an override grant.""" + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = 0 + deployment.litellm_params.allow_client_keepalive_override = True + + router = MagicMock() + router.get_deployment.return_value = deployment + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "deploy-disabled"} + + result = _resolve_keepalive_seconds({"model": "my-model", "keepalive_seconds": 250}, response=response) + assert result == 0.0 + + +def test_keepalive_from_deployment_config_reads_by_model_id(monkeypatch): + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = 45.0 + deployment.litellm_params.allow_client_keepalive_override = True + + router = MagicMock() + router.get_deployment.return_value = deployment + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "deploy-abc"} + + result = _keepalive_from_deployment_config({"model": "my-model"}, response) + assert result == ps._DeploymentKeepaliveConfig(keepalive_seconds=45.0, allow_client_override=True) + router.get_deployment.assert_called_once_with(model_id="deploy-abc") + + +def test_keepalive_from_deployment_config_stale_model_id_does_not_fall_through(monkeypatch): + """A populated model_id names the specific deployment that served the stream. + If that ID no longer resolves (e.g. removed by a config reload mid-stream), + that's a stale identity, not an absent one: it must not fall through to the + model_name fallback, since a currently-live sibling deployment's config was + never what actually served this stream, even if that sibling's config is + unambiguous on its own.""" + from unittest.mock import MagicMock + + router = MagicMock() + router.get_deployment.return_value = None + router.get_model_list.return_value = [ + {"litellm_params": {"keepalive_seconds": 20.0}}, + ] + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {"model_id": "stale-deploy-id"} + + result = _keepalive_from_deployment_config({"model": "slow-model"}, response) + assert result is None + router.get_model_list.assert_not_called() + + +def test_keepalive_from_deployment_config_fallback_by_name(monkeypatch): + from unittest.mock import MagicMock + + router = MagicMock() + router.get_deployment.return_value = None + router.get_model_list.return_value = [ + {"litellm_params": {"keepalive_seconds": 20.0}}, + ] + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {} + + result = _keepalive_from_deployment_config({"model": "slow-model"}, response) + assert result == ps._DeploymentKeepaliveConfig(keepalive_seconds=20.0, allow_client_override=False) + router.get_model_list.assert_called_once_with(model_name="slow-model") + + +def test_keepalive_from_deployment_config_fallback_by_name_agreeing_deployments(monkeypatch): + """Multiple deployments under the same model_name with the same keepalive_seconds + is unambiguous, so the shared value is used even without a model_id.""" + from unittest.mock import MagicMock + + router = MagicMock() + router.get_deployment.return_value = None + router.get_model_list.return_value = [ + {"litellm_params": {"keepalive_seconds": 20.0, "allow_client_keepalive_override": True}}, + {"litellm_params": {"keepalive_seconds": 20.0, "allow_client_keepalive_override": True}}, + ] + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {} + + result = _keepalive_from_deployment_config({"model": "slow-model"}, response) + assert result == ps._DeploymentKeepaliveConfig(keepalive_seconds=20.0, allow_client_override=True) + + +def test_keepalive_from_deployment_config_fallback_by_name_conflicting_deployments(monkeypatch): + """Without a model_id, if deployments under the same model_name disagree on + keepalive_seconds, we can't tell which one served the stream: don't guess and + apply the wrong deployment's interval (or override an explicit disable).""" + from unittest.mock import MagicMock + + router = MagicMock() + router.get_deployment.return_value = None + router.get_model_list.return_value = [ + {"litellm_params": {"keepalive_seconds": 20.0}}, + {"litellm_params": {"keepalive_seconds": 0}}, + ] + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {} + + result = _keepalive_from_deployment_config({"model": "slow-model"}, response) + assert result is None + + +def test_keepalive_from_deployment_config_fallback_by_name_configured_plus_unset(monkeypatch): + """A deployment that leaves keepalive_seconds unset entirely (not explicitly 0) + must not inherit a sibling deployment's configured interval: without a model_id + we can't tell which deployment served the stream, so mixing a configured + deployment with an unconfigured one is just as ambiguous as two conflicting + configured values.""" + from unittest.mock import MagicMock + + router = MagicMock() + router.get_deployment.return_value = None + router.get_model_list.return_value = [ + {"litellm_params": {"keepalive_seconds": 20.0}}, + {"litellm_params": {}}, + ] + + monkeypatch.setattr(ps, "llm_router", router) + + response = MagicMock() + response._hidden_params = {} + + result = _keepalive_from_deployment_config({"model": "slow-model"}, response) + assert result is None + + +def test_keepalive_from_deployment_config_no_router_returns_none(monkeypatch): + monkeypatch.setattr(ps, "llm_router", None) + result = _keepalive_from_deployment_config({"model": "gpt-4"}, None) + assert result is None + + +def test_make_keepalive_resolver_caches_by_model_id(monkeypatch): + """The steady-state case (no fallback): every chunk shares the same + model_id, so the deployment lookup must happen once, not once per chunk.""" + from unittest.mock import MagicMock + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = 5.0 + deployment.litellm_params.allow_client_keepalive_override = False + + router = MagicMock() + router.get_deployment.return_value = deployment + monkeypatch.setattr(ps, "llm_router", router) + + resolve = _make_keepalive_resolver({"model": "my-model"}) + + first = _simple_chunk(content="a") + first._hidden_params = {"model_id": "deploy-steady"} + second = _simple_chunk(content="b") + second._hidden_params = {"model_id": "deploy-steady"} + + assert resolve(first) == 5.0 + assert resolve(second) == 5.0 + router.get_deployment.assert_called_once_with(model_id="deploy-steady") + + +def test_make_keepalive_resolver_reresolves_on_model_id_change(monkeypatch): + """A mid-stream fallback changes model_id: the cache must miss and + re-resolve against the new deployment, not keep serving the stale value.""" + from unittest.mock import MagicMock + + before = MagicMock() + before.litellm_params.keepalive_seconds = 5.0 + before.litellm_params.allow_client_keepalive_override = False + + after = MagicMock() + after.litellm_params.keepalive_seconds = 30.0 + after.litellm_params.allow_client_keepalive_override = False + + router = MagicMock() + router.get_deployment.side_effect = lambda model_id: {"deploy-a": before, "deploy-b": after}[model_id] + monkeypatch.setattr(ps, "llm_router", router) + + resolve = _make_keepalive_resolver({"model": "my-model"}) + + chunk_a = _simple_chunk(content="a") + chunk_a._hidden_params = {"model_id": "deploy-a"} + chunk_b = _simple_chunk(content="b") + chunk_b._hidden_params = {"model_id": "deploy-b"} + + assert resolve(chunk_a) == 5.0 + assert resolve(chunk_b) == 30.0 + assert router.get_deployment.call_count == 2 + + +def test_make_keepalive_resolver_missing_model_id_never_cached(monkeypatch): + """Without a model_id there's no reliable cache key (see the model_name + fallback in _keepalive_from_deployment_config), so every chunk must + re-resolve fresh rather than reuse a stale guess.""" + from unittest.mock import MagicMock + + router = MagicMock() + router.get_deployment.return_value = None + router.get_model_list.return_value = [{"litellm_params": {"keepalive_seconds": 12.0}}] + monkeypatch.setattr(ps, "llm_router", router) + + resolve = _make_keepalive_resolver({"model": "slow-model"}) + + chunk_a = _simple_chunk(content="a") + chunk_a._hidden_params = {} + chunk_b = _simple_chunk(content="b") + chunk_b._hidden_params = {} + + assert resolve(chunk_a) == 12.0 + assert resolve(chunk_b) == 12.0 + assert router.get_model_list.call_count == 2 + + +def test_make_keepalive_resolver_expires_cache_after_ttl(monkeypatch): + """An operator's live config change (revoking override, disabling + keepalive, removing the deployment) must be observed within + _KEEPALIVE_CACHE_TTL_SECONDS, not frozen for the rest of an + already-in-flight stream just because the model_id hasn't changed.""" + from unittest.mock import MagicMock + + before = MagicMock() + before.litellm_params.keepalive_seconds = 20.0 + before.litellm_params.allow_client_keepalive_override = False + + after = MagicMock() + after.litellm_params.keepalive_seconds = 0 + after.litellm_params.allow_client_keepalive_override = False + + router = MagicMock() + router.get_deployment.return_value = before + monkeypatch.setattr(ps, "llm_router", router) + + clock = {"t": 0.0} + monkeypatch.setattr(ps.time, "monotonic", lambda: clock["t"]) + + resolve = _make_keepalive_resolver({"model": "my-model"}) + + chunk = _simple_chunk(content="a") + chunk._hidden_params = {"model_id": "deploy-live"} + + assert resolve(chunk) == 20.0 + assert router.get_deployment.call_count == 1 + + # Still within the TTL: same model_id, cached value reused even though + # the router's live config has since changed underneath it. + router.get_deployment.return_value = after + clock["t"] = ps._KEEPALIVE_CACHE_TTL_SECONDS - 0.01 + assert resolve(chunk) == 20.0 + assert router.get_deployment.call_count == 1 + + # Past the TTL: the config-reload disable is now observed. + clock["t"] = ps._KEEPALIVE_CACHE_TTL_SECONDS + 0.01 + assert resolve(chunk) == 0.0 + assert router.get_deployment.call_count == 2 + + +def test_keepalive_seconds_in_all_litellm_params(): + from litellm.types.utils import all_litellm_params + + assert "keepalive_seconds" in all_litellm_params + + +def test_allow_client_keepalive_override_in_all_litellm_params(): + """allow_client_keepalive_override is a deployment-only control flag: if it's + missing from all_litellm_params, it leaks straight through into the actual + provider API call as an unrecognized field and gets rejected (confirmed live + against the real Anthropic API, which returns 'Extra inputs are not + permitted').""" + from litellm.types.utils import all_litellm_params + + assert "allow_client_keepalive_override" in all_litellm_params + + +@pytest.mark.asyncio +async def test_async_data_generator_emits_ping_heartbeat(monkeypatch): + """When keepalive_seconds is set on a deployment that allows client override, + ': ping' frames appear during upstream stalls.""" + import asyncio + from unittest.mock import MagicMock + + _patch_logging_flags(monkeypatch) + monkeypatch.setattr(ps, "_KEEPALIVE_MIN_SECONDS", 0.05) + + router = MagicMock() + router.get_deployment.return_value = None + router.get_model_list.return_value = [{"litellm_params": {"allow_client_keepalive_override": True}}] + monkeypatch.setattr(ps, "llm_router", router) + + async def _slow_response(): + yield _simple_chunk(content="hello") + await asyncio.sleep(0.4) + yield _simple_chunk(content="world") + + out = [] + async for line in async_data_generator( + response=_slow_response(), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4", "keepalive_seconds": 0.05}, + ): + out.append(line) + + pings = [item for item in out if item == ": ping\n\n"] + assert len(pings) >= 2, f"expected >= 2 ping frames; got {len(pings)}" + assert out[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_async_data_generator_no_keepalive_no_pings(monkeypatch): + """Without keepalive_seconds, no ': ping' frames are emitted.""" + _patch_logging_flags(monkeypatch) + + out = [] + async for line in async_data_generator( + response=_async_iter([_simple_chunk(content="hello")]), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4"}, + ): + out.append(line) + + assert ": ping\n\n" not in out + assert out[-1] == "data: [DONE]\n\n" + + +@pytest.mark.asyncio +async def test_async_data_generator_resolves_deployment_once_per_steady_stream(monkeypatch): + """Regression test for the per-chunk resolver cost: a stream where every + real chunk comes from the same deployment (the common, no-fallback case) + must only pay for one `llm_router.get_deployment()` call, not one per + chunk. Before caching, this asserted 1 but got len(chunks) since the + resolver re-ran the full deployment lookup after every single chunk. + + The very first resolve happens on the raw `response` object before any + chunk is yielded; a bare async generator (unlike the real + CustomStreamWrapper this stands in for) can't carry `_hidden_params`, so + that one call goes through the model_name fallback instead of + `get_deployment` — hence it's asserted separately. + """ + from unittest.mock import MagicMock + + _patch_logging_flags(monkeypatch) + + deployment = MagicMock() + deployment.litellm_params.keepalive_seconds = None + deployment.litellm_params.allow_client_keepalive_override = False + + router = MagicMock() + router.get_deployment.return_value = deployment + router.get_model_list.return_value = [{"litellm_params": {}}] + monkeypatch.setattr(ps, "llm_router", router) + + async def _steady_response(): + for content in ("a", "b", "c", "d", "e"): + chunk = _simple_chunk(content=content) + chunk._hidden_params = {"model_id": "deploy-steady"} + yield chunk + + out = [] + async for line in async_data_generator( + response=_steady_response(), + user_api_key_dict=_user_auth(), + request_data={"model": "gpt-4"}, + ): + out.append(line) + + assert router.get_deployment.call_count == 1 + assert router.get_model_list.call_count == 1 + assert out[-1] == "data: [DONE]\n\n" diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py new file mode 100644 index 00000000000..7f4bd935a2b --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py @@ -0,0 +1,33 @@ +"""Tests for the opt-in flag that gates PTU flat-cost attribution.""" + +import pytest + +from litellm.proxy.spend_tracking.ptu_feature_flag import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + is_ptu_cost_attribution_enabled, +) + + +def test_disabled_when_env_var_is_unset(monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + assert is_ptu_cost_attribution_enabled() is False + + +@pytest.mark.parametrize("value", ["true", "True", "TRUE", " true "]) +def test_enabled_for_the_values_the_house_helper_recognises(monkeypatch, value): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, value) + assert is_ptu_cost_attribution_enabled() is True + + +@pytest.mark.parametrize("value", ["false", "False", "0", "1", "", "yes", "off", "maybe"]) +def test_disabled_for_everything_else(monkeypatch, value): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, value) + assert is_ptu_cost_attribution_enabled() is False + + +def test_reads_the_env_var_on_every_call(monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + assert is_ptu_cost_attribution_enabled() is False + + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + assert is_ptu_cost_attribution_enabled() is True diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py new file mode 100644 index 00000000000..d17f6293cc3 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py @@ -0,0 +1,1504 @@ +"""Tests for the per-model PTU flat-cost daily rollup.""" + +import types +from datetime import date, datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import litellm.proxy.spend_tracking.ptu_flat_cost_rollup as ptu_rollup +from litellm.constants import PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR +from litellm.types.router import ModelInfo +from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( + PTUModel, + _active_hours_on_day, + _compute_daily_flat_cost, + _parse_ptu_model, + run_ptu_flat_cost_backfill, + run_ptu_flat_cost_rollup, + run_scheduled_ptu_rollup, +) + +DAY = date(2026, 7, 30) +TODAY = date(2026, 7, 31) + + +# The endpoints require ptu_effective_from alongside the count and rate, so a fixture that +# omits it would exercise a shape the write path cannot produce. Tests about the start +# itself pass with_start=False. +_DEFAULT_PTU_START = "2020-01-01T00:00:00Z" + + +@pytest.fixture(autouse=True) +def _ptu_enabled(monkeypatch): + """PTU is gated off by default. These cover the rollup's mechanics, not the gate, so + they run with it on; the gate itself is covered by its own test below.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + +_VALID_PTU = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} + + +def _model_row(model_id="m1", model_name="gpt-4o-mini-ptu", model_info=None, with_start=True): + row = MagicMock() + row.model_id = model_id + row.model_name = model_name + if ( + with_start + and isinstance(model_info, dict) + and model_info.get("ptu_count") is not None + and model_info.get("cost_per_ptu_per_hour") is not None + and "ptu_effective_from" not in model_info + ): + model_info = {**model_info, "ptu_effective_from": _DEFAULT_PTU_START} + row.model_info = model_info + return row + + +def _model(**overrides): + base = dict(model_id="m", model_name="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=2.0) + base.update(overrides) + return PTUModel(**base) + + +def test_full_day_when_no_window(): + # 5 PTU * $2.00/hr * 24h = $240 + assert _compute_daily_flat_cost(_model(), DAY) == pytest.approx(240.0) + + +def test_window_opening_at_2300_charges_one_hour(): + m = _model(effective_from=datetime(2026, 7, 30, 23, 0, tzinfo=timezone.utc)) + assert _active_hours_on_day(m, DAY) == pytest.approx(1.0) + # 5 * 2.0 * 1 = 10 + assert _compute_daily_flat_cost(m, DAY) == pytest.approx(10.0) + + +def test_window_closing_at_0600_charges_six_hours(): + m = _model(effective_to=datetime(2026, 7, 30, 6, 0, tzinfo=timezone.utc)) + assert _active_hours_on_day(m, DAY) == pytest.approx(6.0) + assert _compute_daily_flat_cost(m, DAY) == pytest.approx(60.0) + + +def test_window_fully_covering_day_charges_24h(): + m = _model( + effective_from=datetime(2026, 7, 1, tzinfo=timezone.utc), + effective_to=datetime(2026, 8, 1, tzinfo=timezone.utc), + ) + assert _active_hours_on_day(m, DAY) == pytest.approx(24.0) + + +def test_window_before_day_charges_zero(): + m = _model(effective_to=datetime(2026, 7, 29, 12, 0, tzinfo=timezone.utc)) + assert _active_hours_on_day(m, DAY) == 0.0 + assert _compute_daily_flat_cost(m, DAY) == 0.0 + + +def test_window_after_day_charges_zero(): + m = _model(effective_from=datetime(2026, 7, 31, 1, 0, tzinfo=timezone.utc)) + assert _active_hours_on_day(m, DAY) == 0.0 + + +def test_naive_effective_from_is_treated_as_utc(): + parsed = _parse_ptu_model( + _model_row( + model_info={ + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "t", + "ptu_effective_from": "2026-07-30T23:00:00", + } + ) + ) + assert parsed is not None + assert _active_hours_on_day(parsed, DAY) == pytest.approx(1.0) + + +def test_effective_from_with_z_suffix_parses(): + parsed = _parse_ptu_model( + _model_row( + model_info={ + "ptu_count": 1, + "cost_per_ptu_per_hour": 1.0, + "team_id": "t", + "ptu_effective_from": "2026-07-30T18:00:00Z", + } + ) + ) + assert parsed is not None + assert _active_hours_on_day(parsed, DAY) == pytest.approx(6.0) + + +@pytest.mark.parametrize( + "model_info", + [ + None, + {}, + {"ptu_count": 5}, + {"cost_per_ptu_per_hour": 2.0}, + {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0}, # missing team_id + {"ptu_count": 0, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}, + {"ptu_count": 5, "cost_per_ptu_per_hour": -1.0, "team_id": "t"}, + {"ptu_count": "not-int", "cost_per_ptu_per_hour": 2.0, "team_id": "t"}, + ], +) +def test_parse_ptu_model_rejects_invalid(model_info): + assert _parse_ptu_model(_model_row(model_info=model_info)) is None + + +@pytest.fixture(autouse=True) +def _no_retry_backoff(monkeypatch): + """Keep the upsert retry backoff out of the test runtime.""" + monkeypatch.setattr(ptu_rollup, "_UPSERT_RETRY_BACKOFF_SECONDS", 0) + + +def _sentinel_row(row_id, team_id, model): + row = MagicMock() + row.id = row_id + row.team_id = team_id + row.model = model + return row + + +def _prisma_with_models(rows, existing_sentinel_rows=()): + prisma = MagicMock() + model_table = MagicMock() + model_table.find_many = AsyncMock(return_value=rows) + daily = MagicMock() + daily.find_many = AsyncMock(return_value=list(existing_sentinel_rows)) + daily.upsert = AsyncMock() + daily.delete_many = AsyncMock() + prisma.db = types.SimpleNamespace(litellm_proxymodeltable=model_table, litellm_dailyteamspend=daily) + return prisma, daily + + +@pytest.mark.asyncio +async def test_rollup_writes_sentinel_row_with_hourly_cost(): + rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "team_x"})] + prisma, table = _prisma_with_models(rows) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.models_processed == 1 + assert result.rows_written == 1 + created = table.upsert.await_args.kwargs["data"]["create"] + assert created["api_key"] == PTU_SENTINEL_API_KEY + assert created["ptu_flat_cost"] == pytest.approx(240.0) + assert created["team_id"] == "team_x" + # identity in the key, display beside it, so a rename cannot move the row + assert created["model"] == "m1" + assert created["model_group"] == "gpt-4o-mini-ptu" + keyed = table.upsert.await_args.kwargs["where"][ + "team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" + ] + assert keyed["model"] == "m1" + + +@pytest.mark.asyncio +async def test_rollup_prunes_stale_row_when_config_is_gone(): + prisma, table = _prisma_with_models( + [_model_row(model_info={"team_id": "team_x"})], + existing_sentinel_rows=[_sentinel_row("stale-1", "team_x", "gpt-4o-mini-ptu")], + ) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.rows_written == 0 + table.upsert.assert_not_awaited() + table.delete_many.assert_awaited_once() + where = table.delete_many.await_args.kwargs["where"] + assert where["date"] == DAY.isoformat() + assert where["api_key"] == PTU_SENTINEL_API_KEY + # the row is garbage because this run did not refresh it, not because of a key list + assert "lt" in where["updated_at"] + + +@pytest.mark.asyncio +async def test_rollup_writes_current_row_before_pruning_and_keeps_it(): + prisma, table = _prisma_with_models( + [_model_row(model_id="ptu", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "team_x"})], + existing_sentinel_rows=[ + _sentinel_row("live", "team_x", "gpt-4o-mini-ptu"), + _sentinel_row("stale", "team_x", "removed-model"), + ], + ) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.rows_written == 1 + table.upsert.assert_awaited_once() + # the upsert lands before the cutoff is applied, so the refreshed row is out of reach + upsert_order = table.method_calls.index(("upsert", (), table.upsert.call_args.kwargs)) + assert upsert_order < [c[0] for c in table.method_calls].index("delete_many") + + +@pytest.mark.asyncio +async def test_two_deployments_sharing_a_name_get_a_row_each(): + """Keyed on the deployment id they no longer need collapsing, and each keeps its own + amount. The read path merges them back under the shared display name.""" + rows = [ + _model_row(model_id="dep-b", model_info={"ptu_count": 2, "cost_per_ptu_per_hour": 1.0, "team_id": "team_x"}), + _model_row(model_id="dep-a", model_info={"ptu_count": 3, "cost_per_ptu_per_hour": 1.0, "team_id": "team_x"}), + ] + prisma, table = _prisma_with_models(rows) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.rows_written == 2 + written = {c.kwargs["data"]["create"]["model"]: c.kwargs["data"]["create"] for c in table.upsert.await_args_list} + assert set(written) == {"dep-a", "dep-b"} + assert written["dep-a"]["ptu_flat_cost"] == pytest.approx(72.0) + assert written["dep-b"]["ptu_flat_cost"] == pytest.approx(48.0) + assert {row["model_group"] for row in written.values()} == {"gpt-4o-mini-ptu"} + + +@pytest.mark.asyncio +async def test_rollup_skips_zero_active_hours(): + rows = [ + _model_row( + model_info={ + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "team_x", + "ptu_effective_from": "2026-08-01T00:00:00Z", + } + ) + ] + prisma, table = _prisma_with_models(rows) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.models_processed == 1 + assert result.rows_written == 0 + table.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_rollup_skips_models_without_ptu_config(): + rows = [ + _model_row(model_id="plain", model_info={"team_id": "team_x"}), + _model_row(model_id="ptu", model_info={"ptu_count": 3, "cost_per_ptu_per_hour": 1.0, "team_id": "team_y"}), + ] + prisma, table = _prisma_with_models(rows) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.models_processed == 1 + assert result.rows_written == 1 + + +def test_parse_ptu_model_skips_a_deployment_with_no_effective_start(): + """The endpoints require a start. A row without one predates that rule or was written + around them, and inferring a start would bill days the deployment did not exist: before + this, a windowless deployment accrued the whole cap window on its first run.""" + assert ( + _parse_ptu_model( + _model_row( + model_info={"ptu_count": 10, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}, + with_start=False, + ) + ) + is None + ) + + +def test_parse_ptu_model_accepts_json_string_model_info(): + # Some query paths deliver model_info as a JSON string, not a dict. + import json as _json + + raw = _json.dumps( + { + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "team_x", + "ptu_effective_from": _DEFAULT_PTU_START, + } + ) + parsed = _parse_ptu_model(_model_row(model_info=raw)) + assert parsed is not None + assert parsed.ptu_count == 5 and parsed.team_id == "team_x" + + +def test_parse_ptu_model_rejects_unparseable_string(): + assert _parse_ptu_model(_model_row(model_info="not-json")) is None + + +def test_parse_ptu_model_accepts_datetime_object_effective_from(): + # model_info can carry a real datetime object, not just an ISO string. + parsed = _parse_ptu_model( + _model_row( + model_info={ + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "t", + "ptu_effective_from": datetime(2026, 7, 30, 23, 0, tzinfo=timezone.utc), + } + ) + ) + assert parsed is not None + assert _active_hours_on_day(parsed, DAY) == pytest.approx(1.0) + + +@pytest.mark.parametrize( + "bounds", + [ + {"ptu_effective_from": "not-a-date"}, + {"ptu_effective_to": 12345}, + {"ptu_effective_from": "not-a-date", "ptu_effective_to": 12345}, + ], +) +def test_parse_ptu_model_rejects_malformed_effective_dates(bounds): + # Treating an unparseable bound as "no bound" would widen the window to the whole + # day and overcharge, so the deployment is skipped until the config is fixed. + parsed = _parse_ptu_model( + _model_row(model_info={"ptu_count": 2, "cost_per_ptu_per_hour": 1.0, "team_id": "t", **bounds}) + ) + assert parsed is None + + +def test_parse_ptu_model_rejects_an_inverted_window(): + # An end at or before the start can only mean a broken config; charging it as an + # open-ended window would bill a full day. + parsed = _parse_ptu_model( + _model_row( + model_info={ + "ptu_count": 2, + "cost_per_ptu_per_hour": 1.0, + "team_id": "t", + "ptu_effective_from": "2026-07-31T12:00:00Z", + "ptu_effective_to": "2026-07-31T06:00:00Z", + } + ) + ) + assert parsed is None + + +@pytest.mark.asyncio +async def test_rollup_returns_empty_when_prisma_client_is_none(): + result = await run_ptu_flat_cost_rollup(None, target_date=DAY) + assert result.models_processed == 0 + assert result.rows_written == 0 + assert result.day == DAY + + +@pytest.mark.asyncio +async def test_rollup_continues_after_a_failed_upsert(): + rows = [ + _model_row(model_id="a", model_info={"ptu_count": 1, "cost_per_ptu_per_hour": 1.0, "team_id": "team_a"}), + _model_row(model_id="b", model_info={"ptu_count": 2, "cost_per_ptu_per_hour": 1.0, "team_id": "team_b"}), + ] + prisma, table = _prisma_with_models(rows) + table.upsert = AsyncMock(side_effect=RuntimeError("db down")) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + # both models exhausted their retries, and the batch still ran to completion + assert result.models_processed == 2 + assert result.rows_written == 0 + assert result.rows_failed == 2 + assert table.upsert.await_count == 2 * ptu_rollup._UPSERT_ATTEMPTS + + +@pytest.mark.asyncio +async def test_rollup_retries_a_transient_upsert_failure_and_succeeds(): + rows = [_model_row(model_info={"ptu_count": 1, "cost_per_ptu_per_hour": 1.0, "team_id": "team_a"})] + prisma, table = _prisma_with_models(rows) + table.upsert = AsyncMock(side_effect=[RuntimeError("connection reset"), None]) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + # the retry writes the day's charge, so nothing is left for a manual rerun + assert result.rows_written == 1 + assert result.rows_failed == 0 + assert table.upsert.await_count == 2 + + +def _pod_lock(acquired): + """A lock manager that acquires (or not) and, by default, still owns the lease.""" + lock = MagicMock() + lock.pod_id = "this-pod" + lock.redis_cache = MagicMock() + lock.redis_cache.async_get_cache = AsyncMock(return_value="this-pod") + lock.get_redis_lock_key = MagicMock(return_value="lock-key") + lock.acquire_lock = AsyncMock(return_value=acquired) + lock.release_lock = AsyncMock() + return lock + + +@pytest.mark.asyncio +async def test_scheduled_rollup_skips_the_run_when_another_pod_holds_the_lock(): + rows = [_model_row(model_info={"ptu_count": 1, "cost_per_ptu_per_hour": 1.0, "team_id": "team_a"})] + prisma, table = _prisma_with_models(rows) + lock = _pod_lock(acquired=False) + + result = await run_scheduled_ptu_rollup(prisma, pod_lock_manager=lock, target_date=DAY) + + # the losing pod must not write or prune, or it could delete the winner's fresh rows + assert result is None + assert table.upsert.await_count == 0 + assert table.delete_many.await_count == 0 + lock.release_lock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_scheduled_rollup_runs_and_releases_the_lock_when_it_wins(): + rows = [_model_row(model_info={"ptu_count": 1, "cost_per_ptu_per_hour": 1.0, "team_id": "team_a"})] + prisma, table = _prisma_with_models(rows) + lock = _pod_lock(acquired=True) + + result = await run_scheduled_ptu_rollup(prisma, pod_lock_manager=lock, target_date=DAY) + + assert result is not None and result.rows_written == 1 + lock.acquire_lock.assert_awaited_once() + lock.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scheduled_rollup_releases_the_lock_even_when_the_run_raises(): + prisma, table = _prisma_with_models([]) + prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=RuntimeError("db down")) + lock = _pod_lock(acquired=True) + + with pytest.raises(RuntimeError): + await run_scheduled_ptu_rollup(prisma, pod_lock_manager=lock, target_date=DAY) + + # a stuck lock would block every later run until its TTL expires + lock.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_scheduled_rollup_runs_unguarded_without_a_redis_backed_lock(): + rows = [_model_row(model_info={"ptu_count": 1, "cost_per_ptu_per_hour": 1.0, "team_id": "team_a"})] + prisma, table = _prisma_with_models(rows) + lock = _pod_lock(acquired=True) + lock.redis_cache = None + + result = await run_scheduled_ptu_rollup(prisma, pod_lock_manager=lock, target_date=DAY) + + # single-writer deployments have no lock to take, and must still reconcile the day + assert result is not None and result.rows_written == 1 + lock.acquire_lock.assert_not_awaited() + + assert await run_scheduled_ptu_rollup(prisma, target_date=DAY) is not None + + +@pytest.mark.asyncio +async def test_rollup_skips_the_prune_when_a_replacement_write_failed(): + # The deployment was renamed, so the old sentinel row is stale only once its + # replacement lands. Pruning against the intended charges after a failed write + # would delete the old row and leave the team with no charge at all. + prisma, table = _prisma_with_models( + [ + _model_row( + model_name="renamed-ptu", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} + ) + ], + existing_sentinel_rows=[_sentinel_row("previous", "t", "old-name-ptu")], + ) + table.upsert = AsyncMock(side_effect=RuntimeError("db down")) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.rows_failed == 1 + table.delete_many.assert_not_awaited() + + +class _FakeSentinelTable: + """In-memory LiteLLM_DailyTeamSpend that honours the sentinel key and prune predicate.""" + + def __init__(self, upsert_gate=None): + self.rows = {} + self._upsert_gate = upsert_gate + self.upsert_keys = [] + self.delete_many_calls = [] + self.find_many_calls = [] + + async def upsert(self, where, data): + if self._upsert_gate is not None: + await self._upsert_gate.wait() + key = where["team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"] + row_key = (key["team_id"], key["date"], key["api_key"], key["model"]) + self.upsert_keys.append(row_key) + self.rows[row_key] = { + "ptu_flat_cost": data["create"]["ptu_flat_cost"], + "model_group": data["create"]["model_group"], + "updated_at": datetime.now(timezone.utc), + } + + async def delete_many(self, where): + self.delete_many_calls.append(where) + cutoff = where["updated_at"]["lt"] + doomed = [ + k + for k, v in self.rows.items() + if k[1] == where["date"] and k[2] == where["api_key"] and v["updated_at"] < cutoff + ] + for k in doomed: + del self.rows[k] + + async def find_many(self, where=None): + """Read back sentinel rows the way prisma would, honouring api_key and a date range.""" + self.find_many_calls.append(where) + if not where or where.get("api_key") != PTU_SENTINEL_API_KEY: + return [] + bounds = where.get("date") or {} + return [ + _stored_sentinel_row(team_id, day, model_id, value.get("model_group")) + for (team_id, day, api_key, model_id), value in self.rows.items() + if api_key == PTU_SENTINEL_API_KEY and (not bounds or bounds["gte"] <= day <= bounds["lte"]) + ] + + def seed(self, team_id, day, model_id, flat_cost, updated_at=None, model_group=None): + """Seed a row the way the rollup writes one: keyed on the deployment id.""" + self.rows[(team_id, day.isoformat(), PTU_SENTINEL_API_KEY, model_id)] = { + "ptu_flat_cost": flat_cost, + "model_group": model_group or model_id, + "updated_at": updated_at or datetime.now(timezone.utc), + } + + +def _stored_sentinel_row(team_id, day, model_id, model_group=None): + row = MagicMock() + row.team_id = team_id + row.date = day + row.model = model_id + row.model_group = model_group + return row + + +def _prisma_for(model_rows, daily_table): + prisma = MagicMock() + model_table = MagicMock() + model_table.find_many = AsyncMock(return_value=model_rows) + prisma.db = types.SimpleNamespace(litellm_proxymodeltable=model_table, litellm_dailyteamspend=daily_table) + return prisma + + +@pytest.mark.asyncio +async def test_an_older_run_cannot_delete_a_newer_runs_row(): + """The race the absolute predicate exists for: an admin renames a PTU model while two + pods are mid-rollup, so each pod prices a different model name. The pod that started + first must not be able to delete the charge the second pod just wrote.""" + import asyncio + + gate = asyncio.Event() + table = _FakeSentinelTable() + ptu = {"ptu_count": 10, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} + + # pod A read the config before the second deployment appeared and is stalled mid-upsert + slow_table = _FakeSentinelTable(upsert_gate=gate) + slow_table.rows = table.rows + pod_a = asyncio.create_task( + run_ptu_flat_cost_rollup( + _prisma_for([_model_row(model_id="dep-a", model_info=ptu)], slow_table), target_date=DAY + ) + ) + await asyncio.sleep(0) # let pod A capture run_started and reach the gate + + # pod B read a config that has since replaced it, and completes first + await run_ptu_flat_cost_rollup(_prisma_for([_model_row(model_id="dep-b", model_info=ptu)], table), target_date=DAY) + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-b") in table.rows + + gate.set() + await pod_a + + # pod A's cutoff predates every row written during the race, so its delete reaches none + assert table.rows[("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-b")]["ptu_flat_cost"] == pytest.approx(480.0) + + +@pytest.mark.asyncio +async def test_a_later_clean_run_clears_the_row_the_race_left_behind(): + """The race can leave a charge for a since-removed deployment in place for a day; the + next run, seeing only the current config, must sweep it.""" + table = _FakeSentinelTable() + ptu = {"ptu_count": 10, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} + stale_key = ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-removed") + table.rows[stale_key] = { + "ptu_flat_cost": 480.0, + "model_group": "retired", + "updated_at": datetime(2020, 1, 1, tzinfo=timezone.utc), + } + + await run_ptu_flat_cost_rollup( + _prisma_for([_model_row(model_id="dep-live", model_info=ptu)], table), target_date=DAY + ) + + assert stale_key not in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows + + +@pytest.mark.asyncio +async def test_scheduled_rollup_alerts_when_a_team_charge_never_landed(): + """A failed charge is a silent underbill: the team shows no PTU cost for the date and + the next cron run moves on to the next day. It has to reach an operator.""" + rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})] + prisma, table = _prisma_with_models(rows) + table.upsert = AsyncMock(side_effect=RuntimeError("db down")) + alert = AsyncMock() + + result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert) + + assert result.rows_failed == 1 + alert.assert_awaited_once() + message = alert.await_args.args[0] + assert DAY.isoformat() in message + assert "rerun" in message + + +@pytest.mark.asyncio +async def test_scheduled_rollup_stays_quiet_when_every_charge_landed(): + rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})] + prisma, table = _prisma_with_models(rows) + alert = AsyncMock() + + result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert) + + assert result.rows_failed == 0 + alert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_a_broken_alert_channel_does_not_fail_the_rollup(): + rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})] + prisma, table = _prisma_with_models(rows) + table.upsert = AsyncMock(side_effect=RuntimeError("db down")) + + result = await run_scheduled_ptu_rollup( + prisma, target_date=DAY, alert=AsyncMock(side_effect=RuntimeError("slack down")) + ) + + # losing the alert must not also lose the run's result or leave the lock held + assert result.rows_failed == 1 + + +@pytest.mark.asyncio +async def test_scheduled_rollup_alerts_from_under_the_pod_lock_too(): + rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})] + prisma, table = _prisma_with_models(rows) + table.upsert = AsyncMock(side_effect=RuntimeError("db down")) + lock = _pod_lock(acquired=True) + alert = AsyncMock() + + await run_scheduled_ptu_rollup(prisma, pod_lock_manager=lock, target_date=DAY, alert=alert) + + alert.assert_awaited_once() + lock.release_lock.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lock_read", + [ + pytest.param(AsyncMock(side_effect=RuntimeError("redis down")), id="redis-unreachable"), + pytest.param(AsyncMock(return_value=None), id="lock-key-missing"), + ], +) +async def test_scheduled_rollup_runs_the_day_when_the_lock_is_unavailable_but_unheld(lock_read): + """acquire_lock reports contention and a Redis outage identically. Treating both as + "someone else has it" would skip the day on every pod at once, losing every team's + charge for that date; the reconcile is safe to run twice, so the day wins.""" + rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})] + prisma, table = _prisma_with_models(rows) + lock = _pod_lock(acquired=False) + lock.redis_cache.async_get_cache = lock_read + + result = await run_scheduled_ptu_rollup(prisma, pod_lock_manager=lock, target_date=DAY) + + assert result is not None and result.rows_written == 1 + table.upsert.assert_awaited_once() + + +def _team_scoped_row(public_name, model_id="m1", team_id="team_x", **ptu): + """A deployment as POST /model/new actually stores it: synthetic routing name in + model_name, the operator's chosen name in model_info.team_public_model_name.""" + info = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": team_id, **ptu} + info["team_public_model_name"] = public_name + return _model_row(model_id=model_id, model_name=f"model_name_{team_id}_{model_id}-uuid", model_info=info) + + +def test_parse_ptu_model_keys_on_the_public_name_not_the_routing_key(): + # PTU requires a team_id, so every PTU deployment carries the synthetic model_name. + # Keying the charge on it files the cost under a UUID no usage view can resolve. + parsed = _parse_ptu_model(_team_scoped_row("gpt-4o")) + assert parsed is not None + assert parsed.model_name == "gpt-4o" + + +def test_parse_ptu_model_falls_back_to_model_name_without_a_public_name(): + parsed = _parse_ptu_model( + _model_row( + model_name="plain-deployment", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} + ) + ) + assert parsed is not None + assert parsed.model_name == "plain-deployment" + + +@pytest.mark.parametrize("bad_public_name", ["", None, 123, {"nested": "value"}]) +def test_parse_ptu_model_ignores_an_unusable_public_name(bad_public_name): + row = _model_row( + model_name="routing-key", + model_info={ + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "team_id": "t", + "team_public_model_name": bad_public_name, + }, + ) + parsed = _parse_ptu_model(row) + assert parsed is not None + assert parsed.model_name == "routing-key" + + +@pytest.mark.asyncio +async def test_team_scoped_deployments_key_on_their_id_and_display_the_public_name(): + """A team-scoped deployment's model_name is a synthetic routing key, so the row keys on + the stable id and carries the operator-facing name alongside it for display.""" + rows = [ + _team_scoped_row("gpt-4o", model_id="dep-b", ptu_count=2, cost_per_ptu_per_hour=1.0), + _team_scoped_row("gpt-4o", model_id="dep-a", ptu_count=3, cost_per_ptu_per_hour=1.0), + ] + prisma, table = _prisma_with_models(rows) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.rows_written == 2 + written = {c.kwargs["data"]["create"]["model"]: c.kwargs["data"]["create"] for c in table.upsert.await_args_list} + assert set(written) == {"dep-a", "dep-b"} + assert {row["model_group"] for row in written.values()} == {"gpt-4o"} + assert sum(row["ptu_flat_cost"] for row in written.values()) == pytest.approx(120.0) + + +# --------------------------------------------------------------------------- +# Catch-up backfill: the days a once-daily "price yesterday" job never revisits +# --------------------------------------------------------------------------- + + +def _day(offset): + """A UTC date relative to DAY, which is the last day the backfill may price.""" + return DAY + timedelta(days=offset) + + +def _windowed_row(effective_from=None, effective_to=None, **overrides): + ptu = { + "ptu_count": overrides.pop("ptu_count", 5), + "cost_per_ptu_per_hour": overrides.pop("cost_per_ptu_per_hour", 2.0), + "team_id": overrides.pop("team_id", "t"), + } + if effective_from is not None: + ptu["ptu_effective_from"] = effective_from.isoformat() + if effective_to is not None: + ptu["ptu_effective_to"] = effective_to.isoformat() + return _model_row(model_info=ptu, **overrides) + + +def _midnight(day): + return datetime.combine(day, datetime.min.time(), tzinfo=timezone.utc) + + +def _priced_dates(table): + return sorted(key[1] for key in table.rows) + + +# --- R1: the gap is actually closed ---------------------------------------- + + +@pytest.mark.asyncio +async def test_backfill_prices_every_elapsed_in_window_day(): + """The defect this exists for: an operator backdates a PTU window by 30 days, the + config validates and persists, and the once-daily job prices only yesterday. Every + elapsed day inside the declared window has to end up with a charge.""" + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-29)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.days_scanned == 30 + assert result.rows_written == 30 + assert result.rows_failed == 0 + assert _priced_dates(table) == [_day(offset).isoformat() for offset in range(-29, 1)] + assert all(row["ptu_flat_cost"] == pytest.approx(240.0) for row in table.rows.values()) + + +@pytest.mark.asyncio +async def test_backfill_prices_a_day_the_daily_run_missed(): + """A pod restart across 00:15 loses exactly one day. Only that day may be written.""" + table = _FakeSentinelTable() + table.seed("t", _day(-2), "m1", 240.0) + table.seed("t", _day(0), "m1", 240.0) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-2)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.rows_written == 1 + assert table.upsert_keys == [("t", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "m1")] + + +@pytest.mark.asyncio +async def test_backfill_prices_the_partial_first_day_by_active_hours(): + """A backfilled day is priced by the same hourly overlap as a live one, so a window + opening at 08:01 charges the remaining 15h59m rather than a whole day.""" + table = _FakeSentinelTable() + opens_at = datetime(_day(-1).year, _day(-1).month, _day(-1).day, 8, 1, tzinfo=timezone.utc) + prisma = _prisma_for([_windowed_row(effective_from=opens_at)], table) + + await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + first_day = table.rows[("t", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "m1")] + assert first_day["ptu_flat_cost"] == pytest.approx(10 * (15 + 59 / 60)) + assert table.rows[("t", _day(0).isoformat(), PTU_SENTINEL_API_KEY, "m1")]["ptu_flat_cost"] == pytest.approx(240.0) + + +@pytest.mark.asyncio +async def test_backfill_stops_at_yesterday(): + """A day that has not finished cannot be billed, however far into the future the + declared window runs.""" + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-2)), effective_to=_midnight(_day(30)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.end == DAY + assert max(_priced_dates(table)) == DAY.isoformat() + + +# --- R2: history is never rewritten ---------------------------------------- + + +@pytest.mark.asyncio +async def test_backfill_leaves_an_existing_row_untouched_when_config_changed(): + """A priced day keeps the amount it was billed at. Re-pricing it under today's rate + would silently restate a closed day, which is worse than the gap being fixed.""" + table = _FakeSentinelTable() + table.seed("t", _day(-1), "m1", 240.0) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-1)), cost_per_ptu_per_hour=5.0)], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + already_priced = ("t", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "m1") + assert already_priced not in table.upsert_keys + assert table.rows[already_priced]["ptu_flat_cost"] == pytest.approx(240.0) + assert result.rows_written == 1 + assert table.rows[("t", _day(0).isoformat(), PTU_SENTINEL_API_KEY, "m1")]["ptu_flat_cost"] == pytest.approx(600.0) + + +def test_parse_skips_a_count_too_large_to_price(): + """float(ptu_count) on an unbounded int raises OverflowError, which aborted the whole + run rather than skipping the one deployment carrying it.""" + assert _parse_ptu_model(_model_row(model_info={**_VALID_PTU, "ptu_count": 10**400})) is None + + +@pytest.mark.parametrize("rate", ["NaN", "Infinity", "-Infinity"]) +def test_parse_skips_a_non_finite_rate(rate): + """NaN compares False against every bound, so a bare `< 0` check passed it through and + the deployment accrued a flat cost of nan.""" + assert _parse_ptu_model(_model_row(model_info={**_VALID_PTU, "cost_per_ptu_per_hour": rate})) is None + + +def test_parse_still_accepts_config_at_the_bounds(): + parsed = _parse_ptu_model( + _model_row( + model_info={ + **_VALID_PTU, + "ptu_count": ModelInfo.MAX_PTU_COUNT, + "cost_per_ptu_per_hour": ModelInfo.MAX_COST_PER_PTU_PER_HOUR, + } + ) + ) + assert parsed is not None and parsed.ptu_count == ModelInfo.MAX_PTU_COUNT + + +@pytest.mark.asyncio +async def test_a_bad_row_does_not_abort_pricing_for_other_teams(): + """One unusable deployment must not take the whole day's rollup down with it.""" + table = _FakeSentinelTable() + prisma = _prisma_for( + [ + _model_row(model_id="bad", model_info={**_VALID_PTU, "ptu_count": 10**400}), + _model_row(model_id="good", model_info=_VALID_PTU), + ], + table, + ) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.rows_written == 1 + assert result.models_processed == 1 + + +@pytest.mark.asyncio +async def test_backfill_keeps_the_history_of_a_deployment_that_was_removed(): + """Deleting a deployment stops it accruing, it does not unbill the days it ran. The + backfill deletes nothing, so a closed day survives its deployment.""" + table = _FakeSentinelTable() + table.seed("t", _day(-1), "dep-gone", 480.0, model_group="gone-model") + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-1)), model_id="dep-live")], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert table.rows[("t", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "dep-gone")]["ptu_flat_cost"] == 480.0 + assert table.delete_many_calls == [] + assert ("t", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows + assert result.rows_written == 2 + + +@pytest.mark.asyncio +async def test_backfill_keeps_history_after_every_ptu_deployment_is_gone(): + """With no PTU config left there is nothing to price, and nothing to delete either.""" + table = _FakeSentinelTable() + for offset in (-2, -1): + table.seed("t", _day(offset), "dep-gone", 480.0, model_group="gone-model") + prisma = _prisma_for([_model_row(model_info={"base_model": "gpt-4o"})], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert len(table.rows) == 2 + assert table.delete_many_calls == [] + assert result.rows_written == 0 + + +@pytest.mark.asyncio +async def test_backfill_keeps_history_when_a_window_is_narrowed(): + """Editing an effective window cannot rewrite a bill that was already correct.""" + table = _FakeSentinelTable() + table.seed("t", _day(-3), "dep-1", 480.0, model_group="ptu-a") + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-1)), model_id="dep-1")], table) + + await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert ("t", _day(-3).isoformat(), PTU_SENTINEL_API_KEY, "dep-1") in table.rows + assert table.delete_many_calls == [] + + +@pytest.mark.asyncio +async def test_backfill_never_prunes_by_timestamp(): + """Retirement is by identity. The catch-up must never take the single-day path's + timestamp predicate, which needs the lock and agreeing clocks to be safe.""" + table = _FakeSentinelTable() + table.seed("t", _day(-3), "dep-1", 1.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-3)), model_id="dep-1")], table) + + await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert all("updated_at" not in (call or {}) for call in table.delete_many_calls) + assert ("t", _day(-3).isoformat(), PTU_SENTINEL_API_KEY, "dep-1") in table.rows + + +# --- R3: a gap is per (team, model, date), not per date --------------------- + + +@pytest.mark.asyncio +async def test_backfill_fills_a_second_model_on_a_day_that_already_has_a_row(): + """A day is not covered just because something was priced on it.""" + table = _FakeSentinelTable() + table.seed("t", _day(-1), "a", 240.0) + prisma = _prisma_for( + [ + _windowed_row(effective_from=_midnight(_day(-1)), model_name="model-a", model_id="a"), + _windowed_row(effective_from=_midnight(_day(-1)), model_name="model-b", model_id="b"), + ], + table, + ) + + await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert ("t", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "b") in table.upsert_keys + assert ("t", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "a") not in table.upsert_keys + + +@pytest.mark.asyncio +async def test_backfill_fills_a_second_team_on_a_day_that_already_has_a_row(): + table = _FakeSentinelTable() + table.seed("team-1", _day(-1), "a", 240.0) + prisma = _prisma_for( + [ + _windowed_row(effective_from=_midnight(_day(-1)), model_name="shared-name", model_id="a", team_id="team-1"), + _windowed_row(effective_from=_midnight(_day(-1)), model_name="shared-name", model_id="b", team_id="team-2"), + ], + table, + ) + + await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert ("team-2", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "b") in table.upsert_keys + assert ("team-1", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "a") not in table.upsert_keys + + +@pytest.mark.asyncio +async def test_backfill_keys_gaps_on_the_public_model_name(): + """Sentinel rows are written under the public name, so a gap check reading the + synthetic routing key would never match one and would rewrite it on every run.""" + table = _FakeSentinelTable() + table.seed("team_x", _day(-1), "m1", 240.0) + row = _team_scoped_row( + "gpt-4o", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=_midnight(_day(-1)).isoformat(), + ) + prisma = _prisma_for([row], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert ("team_x", _day(-1).isoformat(), PTU_SENTINEL_API_KEY, "m1") not in table.upsert_keys + assert result.rows_written == 1 + + +# --- R4: bounds and convergence -------------------------------------------- + + +@pytest.mark.asyncio +async def test_backfill_does_not_scan_before_effective_from(): + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-2)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.start == _day(-2) + assert result.days_scanned == 3 + + +@pytest.mark.asyncio +async def test_backfill_caps_lookback_for_a_model_with_no_effective_from(): + """An open-ended window would otherwise scan back to the beginning of the table.""" + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row()], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.start == DAY - timedelta(days=PTU_ROLLUP_MAX_BACKFILL_DAYS) + assert result.days_scanned == PTU_ROLLUP_MAX_BACKFILL_DAYS + 1 + + +@pytest.mark.asyncio +async def test_backfill_caps_lookback_for_a_window_older_than_the_cap(): + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-400)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.start == DAY - timedelta(days=PTU_ROLLUP_MAX_BACKFILL_DAYS) + + +@pytest.mark.asyncio +async def test_backfill_writes_nothing_for_out_of_window_days(): + """A zero-cost day must write no row, or the gap check would read it as priced.""" + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-2)), effective_to=_midnight(_day(-1)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.days_scanned == 3 + assert result.rows_written == 1 + assert _priced_dates(table) == [_day(-2).isoformat()] + + +@pytest.mark.asyncio +async def test_backfill_is_a_no_op_on_a_fully_priced_range(): + table = _FakeSentinelTable() + for offset in (-2, -1, 0): + table.seed("t", _day(offset), "m1", 240.0) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-2)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.rows_written == 0 + assert table.upsert_keys == [] + assert table.delete_many_calls == [] + + +@pytest.mark.asyncio +async def test_backfill_run_twice_is_idempotent(): + """The second pass must be free, including leaving updated_at alone.""" + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-3)))], table) + + await run_ptu_flat_cost_backfill(prisma, today=TODAY) + snapshot = {key: dict(value) for key, value in table.rows.items()} + table.upsert_keys.clear() + + second = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert second.rows_written == 0 + assert table.upsert_keys == [] + assert table.rows == snapshot + + +@pytest.mark.asyncio +async def test_backfill_does_nothing_without_ptu_config(): + table = _FakeSentinelTable() + prisma = _prisma_for([_model_row(model_info={"base_model": "gpt-4o"})], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.days_scanned == 0 + assert result.rows_written == 0 + assert table.upsert_keys == [] + + +@pytest.mark.asyncio +async def test_backfill_writes_nothing_for_a_window_that_opens_tomorrow(): + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(5)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.days_scanned == 0 + assert table.upsert_keys == [] + + +@pytest.mark.asyncio +async def test_backfill_returns_empty_when_prisma_client_is_none(): + result = await run_ptu_flat_cost_backfill(None, today=TODAY) + + assert result.rows_written == 0 + assert result.days_scanned == 0 + + +@pytest.mark.asyncio +async def test_backfill_counts_a_charge_that_never_landed(): + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-1)))], table) + prisma.db.litellm_dailyteamspend.upsert = AsyncMock(side_effect=RuntimeError("db down")) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.rows_written == 0 + assert result.rows_failed == 2 + + +# --- R5: interaction with the daily path ----------------------------------- + + +@pytest.mark.asyncio +async def test_scheduled_rollup_backfills_after_pricing_the_day(): + """The catch-up pass runs after the day's own rollup, so it sees yesterday already + priced and does not write it a second time.""" + table = _FakeSentinelTable() + yesterday = datetime.now(timezone.utc).date() - timedelta(days=1) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(yesterday - timedelta(days=2)))], table) + + await run_scheduled_ptu_rollup(prisma) + + yesterday_key = ("t", yesterday.isoformat(), PTU_SENTINEL_API_KEY, "m1") + assert table.upsert_keys.count(yesterday_key) == 1 + assert len(table.rows) == 3 + + +@pytest.mark.asyncio +async def test_scheduled_rollup_with_an_explicit_target_date_does_not_backfill(): + """An explicit date means reconcile exactly that day, so no catch-up pass runs.""" + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-10)))], table) + + await run_scheduled_ptu_rollup(prisma, target_date=DAY) + + assert _priced_dates(table) == [DAY.isoformat()] + assert table.find_many_calls == [] + + +@pytest.mark.asyncio +async def test_scheduled_rollup_holds_one_lock_across_both_phases(): + """Backfill running outside the lock would let another pod's prune race its writes.""" + table = _FakeSentinelTable() + yesterday = datetime.now(timezone.utc).date() - timedelta(days=1) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(yesterday - timedelta(days=3)))], table) + rows_at_release = [] + lock = _pod_lock(acquired=True) + lock.release_lock = AsyncMock(side_effect=lambda **kwargs: rows_at_release.append(len(table.rows))) + + await run_scheduled_ptu_rollup(prisma, pod_lock_manager=lock) + + lock.acquire_lock.assert_awaited_once() + assert rows_at_release == [4] + + +@pytest.mark.asyncio +async def test_a_failing_backfill_does_not_lose_the_days_rollup_result(): + """The day's rollup has already run and committed; a broken catch-up pass must not + swallow its result or raise into the scheduler.""" + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row()], table) + prisma.db.litellm_dailyteamspend.find_many = AsyncMock(side_effect=RuntimeError("read replica down")) + + result = await run_scheduled_ptu_rollup(prisma) + + assert result is not None + assert result.rows_written == 1 + assert result.rows_failed == 0 + + +@pytest.mark.asyncio +async def test_scheduled_rollup_alerts_when_a_backfill_charge_never_landed(): + """An unpriced day that stays unpriced is the silent underbill this work exists to + remove, so it has to reach an operator too.""" + table = _FakeSentinelTable() + yesterday = datetime.now(timezone.utc).date() - timedelta(days=1) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(yesterday - timedelta(days=1)))], table) + prisma.db.litellm_dailyteamspend.upsert = AsyncMock(side_effect=RuntimeError("db down")) + alert = AsyncMock() + + await run_scheduled_ptu_rollup(prisma, alert=alert) + + messages = [call.args[0] for call in alert.await_args_list] + assert any("backfill" in message for message in messages) + assert any("unpriced" in message for message in messages) + + +@pytest.mark.asyncio +async def test_a_broken_alert_channel_does_not_fail_the_backfill(): + table = _FakeSentinelTable() + prisma = _prisma_for([_windowed_row()], table) + prisma.db.litellm_dailyteamspend.upsert = AsyncMock(side_effect=RuntimeError("db down")) + + result = await run_scheduled_ptu_rollup(prisma, alert=AsyncMock(side_effect=RuntimeError("slack down"))) + + assert result.rows_failed == 1 + + +# --- R6: the shape the cron actually calls --------------------------------- + + +@pytest.mark.asyncio +async def test_scheduled_rollup_with_no_target_date_closes_a_backdated_window(): + """The production call shape from proxy_server.py, on the real clock: no target_date, + a window backdated 30 days, and every elapsed in-window day has to end up priced with + no operator alert raised. Every other rollup test pins target_date, which is exactly + why this regression shipped.""" + table = _FakeSentinelTable() + today = datetime.now(timezone.utc).date() + opened_on = today - timedelta(days=30) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(opened_on))], table) + alert = AsyncMock() + + await run_scheduled_ptu_rollup(prisma, alert=alert) + + expected = [(opened_on + timedelta(days=offset)).isoformat() for offset in range(30)] + assert _priced_dates(table) == expected + alert.assert_not_awaited() + + +# --- R8: a rename must not re-price history under the new name ---------------- + + +@pytest.mark.asyncio +async def test_backfill_does_not_double_price_a_day_after_a_rename(): + """A rename does not move the row, because the key is the deployment id. Every already + priced day stays a single charge and only the unpriced day is written.""" + table = _FakeSentinelTable() + for offset in (-2, -1): + table.seed("t", _day(offset), "dep-1", 240.0, model_group="old-name") + prisma = _prisma_for( + [_windowed_row(effective_from=_midnight(_day(-2)), model_name="new-name", model_id="dep-1")], table + ) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + priced_on = [key for key in table.rows if key[1] == _day(-1).isoformat()] + assert len(priced_on) == 1, f"day {_day(-1)} carries two charges: {priced_on}" + assert result.rows_written == 1 + assert table.upsert_keys == [("t", _day(0).isoformat(), PTU_SENTINEL_API_KEY, "dep-1")] + + +@pytest.mark.asyncio +async def test_backfill_still_prices_a_genuinely_missing_day_for_a_renamed_deployment(): + """Rename safety must not swallow real gaps: a day with no row for the deployment at + all still gets one, and it carries the current display name.""" + table = _FakeSentinelTable() + table.seed("t", _day(-2), "dep-1", 240.0, model_group="old-name") + prisma = _prisma_for( + [_windowed_row(effective_from=_midnight(_day(-2)), model_name="new-name", model_id="dep-1")], table + ) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.rows_written == 2 + assert sorted(key[1] for key in table.upsert_keys) == [_day(-1).isoformat(), _day(0).isoformat()] + assert all(key[3] == "dep-1" for key in table.upsert_keys) + assert table.rows[("t", _day(0).isoformat(), PTU_SENTINEL_API_KEY, "dep-1")]["model_group"] == "new-name" + + +@pytest.mark.asyncio +async def test_backfill_falls_back_to_the_name_when_a_row_carries_no_source_model_id(): + """A sentinel row whose display name is missing still counts as priced: identity is the + model column, so the gap check never depends on the name being present.""" + table = _FakeSentinelTable() + table.seed("t", _day(-1), "m1", 240.0) + prisma = _prisma_for([_windowed_row(effective_from=_midnight(_day(-1)))], table) + + result = await run_ptu_flat_cost_backfill(prisma, today=TODAY) + + assert result.rows_written == 1 + assert table.upsert_keys == [("t", _day(0).isoformat(), PTU_SENTINEL_API_KEY, "m1")] + + +@pytest.mark.asyncio +async def test_two_runs_straddling_a_rename_collapse_onto_one_row(): + """The reported defect. Two runs holding different config views of the same deployment + used to write two keys for one day; keyed on the id they write the same key, so the + upsert collapses them instead of double charging.""" + table = _FakeSentinelTable() + before = _prisma_for( + [_windowed_row(effective_from=_midnight(_day(-1)), model_name="old-name", model_id="dep-1")], table + ) + after = _prisma_for( + [_windowed_row(effective_from=_midnight(_day(-1)), model_name="new-name", model_id="dep-1")], table + ) + + await run_ptu_flat_cost_backfill(before, today=TODAY) + await run_ptu_flat_cost_backfill(after, today=TODAY) + + for offset in (-1, 0): + charges = [key for key in table.rows if key[1] == _day(offset).isoformat()] + assert len(charges) == 1, f"day {_day(offset)} carries {len(charges)} charges: {charges}" + assert sum(row["ptu_flat_cost"] for row in table.rows.values()) == pytest.approx(480.0) + + +# --- R2: the interleaving that reproduced live on a four-pod rig ------------- + + +@pytest.mark.asyncio +async def test_concurrent_runs_straddling_a_rename_write_one_row(): + """The exact shape reproduced on a live multi-pod rig, which double charged a day. + + Both pods read the day as unpriced before either writes, and a rename lands between + their config reads. Keyed on the display name they produced two different composite + keys and both rows survived, permanently. Keyed on the deployment id they produce the + same key, so the upsert collapses them. + """ + import asyncio + + table = _FakeSentinelTable() + ptu = {"ptu_count": 10, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} + read_by_both = asyncio.Event() + + real_keys = ptu_rollup._existing_sentinel_keys + arrivals = [] + + async def gated_keys(*args, **kwargs): + """Hold the first caller until the second has also read, so neither sees the other.""" + keys = await real_keys(*args, **kwargs) + arrivals.append(1) + if len(arrivals) >= 2: + read_by_both.set() + await read_by_both.wait() + return keys + + ptu_rollup._existing_sentinel_keys = gated_keys + try: + pod_a = asyncio.create_task( + run_ptu_flat_cost_backfill( + _prisma_for([_model_row(model_id="dep-1", model_name="old-name", model_info=ptu)], table), + today=TODAY, + ) + ) + pod_b = asyncio.create_task( + run_ptu_flat_cost_backfill( + _prisma_for([_model_row(model_id="dep-1", model_name="new-name", model_info=ptu)], table), + today=TODAY, + ) + ) + await asyncio.wait_for(asyncio.gather(pod_a, pod_b), timeout=5) + finally: + ptu_rollup._existing_sentinel_keys = real_keys + + for day, count in sorted((key[1], 1) for key in table.rows): + assert count == 1 + per_day = {} + for team_id, day, api_key, model_id in table.rows: + per_day[day] = per_day.get(day, 0) + 1 + assert set(per_day.values()) == {1}, f"a day carries more than one charge: {per_day}" + assert {key[3] for key in table.rows} == {"dep-1"} + + +@pytest.mark.asyncio +async def test_a_rate_change_between_concurrent_runs_leaves_one_row(): + """Two pods disagreeing on the rate, not just the name, still land on one row. Last + writer wins on the amount, which is self-consistent rather than a second charge.""" + table = _FakeSentinelTable() + cheap = _prisma_for( + [_windowed_row(effective_from=_midnight(_day(-1)), model_id="dep-1", cost_per_ptu_per_hour=2.0)], table + ) + dear = _prisma_for( + [_windowed_row(effective_from=_midnight(_day(-1)), model_id="dep-1", cost_per_ptu_per_hour=4.0)], table + ) + + await run_ptu_flat_cost_backfill(cheap, today=TODAY) + await run_ptu_flat_cost_backfill(dear, today=TODAY) + + assert len(table.rows) == 2 # one per elapsed in-window day, not per config view + assert all(row["ptu_flat_cost"] == pytest.approx(240.0) for row in table.rows.values()) + + +# --- R6: the prune is the one operation that needs the lock ------------------- + + +@pytest.mark.asyncio +async def test_an_unguarded_run_writes_but_does_not_prune(): + """Without the cross-pod lock the upserts still run, since they are idempotent, but the + delete does not: its cutoff and the rows' updated_at come from different hosts, so a pod + whose clock runs ahead would sweep a charge a concurrent pod just wrote.""" + table = _FakeSentinelTable() + table.seed("t", DAY, "dep-gone", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) + prisma = _prisma_for( + [_model_row(model_id="dep-live", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})], + table, + ) + + await run_scheduled_ptu_rollup(prisma, target_date=DAY) + + assert table.delete_many_calls == [] + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-gone") in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows + + +@pytest.mark.asyncio +async def test_a_run_holding_the_lock_still_prunes(): + """Losing the sweep entirely would leave stale charges forever, so the guarded path, + which is the normal one, keeps it.""" + table = _FakeSentinelTable() + table.seed("t", DAY, "dep-gone", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) + prisma = _prisma_for( + [_model_row(model_id="dep-live", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})], + table, + ) + + await run_scheduled_ptu_rollup(prisma, pod_lock_manager=_pod_lock(acquired=True), target_date=DAY) + + assert table.delete_many_calls != [] + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-gone") not in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows + + +@pytest.mark.asyncio +async def test_the_prune_cutoff_allows_for_clock_skew_between_hosts(): + """A row written seconds ago by a pod whose clock lags must survive; one written hours + ago by a previous run must not. The grace separates the two populations without + requiring the hosts' clocks to agree.""" + table = _FakeSentinelTable() + just_written = datetime.now(timezone.utc) - timedelta(seconds=30) + table.seed("t", DAY, "dep-concurrent", 480.0, updated_at=just_written) + table.seed("t", DAY, "dep-stale", 480.0, updated_at=datetime.now(timezone.utc) - timedelta(hours=6)) + prisma = _prisma_for([], table) + + await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-concurrent") in table.rows, ( + "a charge written 30s ago by a lagging pod was swept" + ) + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-stale") not in table.rows + + +@pytest.mark.asyncio +async def test_scheduled_rollup_writes_nothing_when_ptu_attribution_is_disabled(monkeypatch): + """Startup already skips scheduling the cron, so this guards the function itself: a + deployment that never opted in accrues nothing whatever route reaches the rollup.""" + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + table = _FakeSentinelTable() + prisma = _prisma_for([_model_row(model_info=_VALID_PTU)], table) + + result = await run_scheduled_ptu_rollup(prisma, pod_lock_manager=None, alert=None) + + assert result is None + assert table.rows == {} + assert table.upsert_keys == [] diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index dc4f860ce00..1435547c434 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -6,13 +6,13 @@ sys.path.insert(0, os.path.abspath("../../../..")) import pytest import litellm -from litellm.router import Router from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.proxy.spend_tracking.savings import ( _baseline_usage, compute_autorouter_savings, compute_savings_spend, ) +from litellm.router import Router from litellm.types.utils import Usage @@ -84,6 +84,236 @@ def test_prompt_caching_savings_priced_at_input_minus_cache_read(): assert result.compression == 0.0 +def _net_caching_savings_against_biller(usage_object: dict, model: str = "claude-sonnet-5") -> float: + """True net caching savings, priced by the real cost calculator. + + Bills the request as it happened, then bills the same token total with nothing + cached, and returns the difference. Deriving the expectation from + ``generic_cost_per_token`` rather than restating the formula is what makes these + tests able to fail: a wrong formula in savings.py cannot also be wrong here. + """ + prompt_tokens = usage_object["prompt_tokens"] + uncached = { + "prompt_tokens": prompt_tokens, + "completion_tokens": usage_object["completion_tokens"], + "total_tokens": prompt_tokens + usage_object["completion_tokens"], + "prompt_tokens_details": {"cached_tokens": 0, "cache_creation_tokens": 0, "text_tokens": prompt_tokens}, + } + return _cost_on(model, uncached) - _cost_on(model, usage_object) + + +def _caching_usage(read: int, written: int, text: int = 10, out: int = 100) -> dict: + prompt_tokens = text + read + written + return { + "prompt_tokens": prompt_tokens, + "completion_tokens": out, + "total_tokens": prompt_tokens + out, + "prompt_tokens_details": { + "cached_tokens": read, + "cache_creation_tokens": written, + "text_tokens": text, + }, + "cache_creation_input_tokens": written, + "cache_read_input_tokens": read, + } + + +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") + _, _, cache_write_cost = _flat_rates("claude-sonnet-5") + # Anthropic charges a premium to write; without it this test asserts nothing. + assert cache_write_cost > input_cost + usage_object = _caching_usage(read=20000, written=500) + result = compute_savings_spend( + model="claude-sonnet-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object=usage_object, + ) + assert result.prompt_caching == pytest.approx(_net_caching_savings_against_biller(usage_object)) + # Strictly less than the gross read discount, which is what shipped before. + assert result.prompt_caching < 20000 * (input_cost - cache_read_cost) + assert result.prompt_caching > 0 + + +def test_prompt_caching_savings_go_negative_on_a_write_only_request(): + """A cold turn that writes cache and gets no hits genuinely cost more than not caching.""" + usage_object = _caching_usage(read=0, written=20000) + result = compute_savings_spend( + model="claude-sonnet-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object=usage_object, + ) + true_savings = _net_caching_savings_against_biller(usage_object) + assert true_savings < 0 + assert result.prompt_caching == pytest.approx(true_savings) + assert result.prompt_caching < 0 + + +def test_prompt_caching_savings_negative_when_writes_outweigh_reads(): + """The wrong-sign case: a few hits against a big write bill is still a net loss.""" + usage_object = _caching_usage(read=1000, written=20000) + result = compute_savings_spend( + model="claude-sonnet-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object=usage_object, + ) + true_savings = _net_caching_savings_against_biller(usage_object) + assert true_savings < 0 + assert result.prompt_caching == pytest.approx(true_savings) + # The gross formula reported this as a saving; the sign itself is the regression. + assert result.prompt_caching < 0 + + +def test_read_only_request_is_unchanged_by_the_write_premium(): + """No cache writes means nothing to net out, so the read discount stands alone.""" + input_cost, cache_read_cost = _anthropic_costs("claude-sonnet-5") + result = compute_savings_spend( + model="claude-sonnet-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object=_caching_usage(read=20000, written=0), + ) + assert result.prompt_caching == pytest.approx(20000 * (input_cost - cache_read_cost)) + + +def test_openai_style_cache_write_tokens_are_netted_out(): + """Providers reporting writes under prompt_tokens_details are netted the same way.""" + _, _, cache_write_cost = _flat_rates("claude-sonnet-5") + input_cost, _ = _anthropic_costs("claude-sonnet-5") + with_top_level = compute_savings_spend( + model="claude-sonnet-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object={"cache_read_input_tokens": 5000, "cache_creation_input_tokens": 800}, + ) + nested_only = compute_savings_spend( + model="claude-sonnet-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object={ + "prompt_tokens_details": {"cached_tokens": 5000, "cache_write_tokens": 800}, + }, + ) + assert nested_only.prompt_caching == pytest.approx(with_top_level.prompt_caching) + assert nested_only.prompt_caching == pytest.approx( + 5000 * (input_cost - _anthropic_costs("claude-sonnet-5")[1]) - 800 * (cache_write_cost - input_cost) + ) + + +def test_model_without_a_cache_write_price_takes_no_premium(): + """An absent write price must mean zero premium, never a bonus. + + ``_get_cost_per_unit`` in the cost calculator defaults a missing price to 0.0. Were + that default copied here the premium would be ``0 - input_cost``, and a model with no + write pricing would report cache writes as free money. This is the common case: most + of the pricing map publishes a cache-read price and no cache-write price. + """ + model = "amazon.nova-2-lite-v1:0" + info = litellm.get_model_info(model=model) + input_cost = info["input_cost_per_token"] + cache_read_cost = info["cache_read_input_token_cost"] + assert info.get("cache_creation_input_token_cost") is None, ( + "fixture drifted: this test needs a model that publishes no cache-write price" + ) + + result = compute_savings_spend( + model=model, + custom_llm_provider=None, + compression_saved_tokens=0, + usage_object=_caching_usage(read=5000, written=5000), + ) + assert result.prompt_caching == pytest.approx(5000 * (input_cost - cache_read_cost)) + assert result.prompt_caching > 0 + + +def test_zero_cache_write_price_is_read_as_unpublished(): + """A ``0.0`` write price means "no separate price", not "writes are free". + + ``deepseek-chat`` carries an explicit zero in the pricing map. Taken literally the + premium would be ``0 - input_cost``, paying out a saving of ``writes * input_cost`` + on traffic that cached nothing. No provider gives cache writes away, so a falsy + price falls open to the input cost like an absent one does. + """ + info = litellm.get_model_info(model="deepseek-chat", custom_llm_provider="deepseek") + assert info.get("cache_creation_input_token_cost") == 0.0, ( + "fixture drifted: this test exists because deepseek-chat publishes a literal 0.0 write price" + ) + + result = compute_savings_spend( + model="deepseek-chat", + custom_llm_provider="deepseek", + compression_saved_tokens=0, + usage_object=_caching_usage(read=0, written=10000), + ) + assert result.prompt_caching == pytest.approx(0.0) + + +def test_zero_cache_read_price_stays_literal(): + """The read leg must NOT copy the write leg's falsy fall-open. + + The two zeros mean opposite things. A free cache *write* is unpublished pricing, so + it falls open to input. A free cache *read* is real and is the largest discount + available -- 15 models charge for input and serve reads for nothing. Falling that + open to the input cost would zero out their savings entirely. + """ + model = "gemini-robotics-er-1.5-preview" + info = litellm.get_model_info(model=model) + input_cost = info["input_cost_per_token"] + assert info.get("cache_read_input_token_cost") == 0.0 and input_cost > 0, ( + "fixture drifted: this test needs a model with paid input and free cache reads" + ) + + result = compute_savings_spend( + model=model, + custom_llm_provider=None, + compression_saved_tokens=0, + usage_object=_caching_usage(read=10000, written=0), + ) + # free reads => the whole input rate is saved, not zero + assert result.prompt_caching == pytest.approx(10000 * input_cost) + + +def test_sub_input_cache_write_price_is_an_extra_saving(): + """A few models price writes below input; there the premium is a real credit. + + Clamping the premium at zero would silently undercount these, so the subtraction + stays signed. ``azure/eu/gpt-4o-2024-11-20`` ships a write price at ~0.5x input. + """ + model = "azure/eu/gpt-4o-2024-11-20" + info = litellm.get_model_info(model=model) + input_cost = info["input_cost_per_token"] + cheap_write = info["cache_creation_input_token_cost"] + assert 0 < cheap_write < input_cost, "fixture drifted: this test needs a model pricing cache writes below input" + # no published read price, so the read leg mirrors input and contributes nothing; + # the whole result is the negative premium, i.e. a credit. + assert info.get("cache_read_input_token_cost") is None + + result = compute_savings_spend( + model=model, + custom_llm_provider=None, + compression_saved_tokens=0, + usage_object=_caching_usage(read=1000, written=4000), + ) + assert result.prompt_caching == pytest.approx(4000 * (input_cost - cheap_write)) + assert result.prompt_caching > 0 + + +def test_negative_cache_write_count_clamps_to_zero(): + """A malformed negative write count must not be read as a saving.""" + input_cost, cache_read_cost = _anthropic_costs("claude-sonnet-5") + result = compute_savings_spend( + model="claude-sonnet-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object={"cache_read_input_tokens": 1000, "cache_creation_input_tokens": -5000}, + ) + assert result.prompt_caching == pytest.approx(1000 * (input_cost - cache_read_cost)) + + def test_unknown_model_fails_open_to_zero(): result = compute_savings_spend( model="totally-made-up-model-xyz", @@ -664,6 +894,47 @@ def test_a_non_string_recorded_baseline_is_ignored(): assert result.autorouter == 0.0 +def test_prompt_caching_prices_at_the_deployment_rate_not_the_public_one(): + """A deployment's negotiated cache rates are what it really pays. + + Pricing the write premium off the public map instead reports a loss ~3x the real + one here, which is the whole point of resolving deployment pricing first. + """ + router = Router( + model_list=[ + { + "model_name": "cheap-sonnet", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5", + "input_cost_per_token": 1e-06, + "cache_creation_input_token_cost": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + }, + }, + ] + ) + deployment_id = router.get_model_list(model_name="cheap-sonnet")[0]["model_info"]["id"] + + result = compute_savings_spend( + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object=_caching_usage(read=1000, written=20000), + model_id=deployment_id, + llm_router=lambda: router, + ) + at_deployment_rates = 1000 * (1e-06 - 1e-07) - 20000 * (1.25e-06 - 1e-06) + assert result.prompt_caching == pytest.approx(at_deployment_rates) + + at_public_rates = compute_savings_spend( + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + usage_object=_caching_usage(read=1000, written=20000), + ) + assert result.prompt_caching > at_public_rates.prompt_caching + + 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_lazy_openapi_snapshot.py b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py index 64cb931888b..79330b0e3a6 100644 --- a/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py +++ b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py @@ -1,6 +1,7 @@ import sys from types import ModuleType, SimpleNamespace +from litellm.proxy._lazy_features import LazyFeature from litellm.proxy._lazy_openapi_snapshot import _normalize_operation_ids @@ -22,22 +23,20 @@ def test_generate_snapshot_uses_shared_operation_id_reservations(monkeypatch): fake_lazy_features_module = ModuleType("litellm.proxy._lazy_features") fake_lazy_features_module.LAZY_FEATURES = [ - SimpleNamespace( + LazyFeature( name="feature-a", module_path="fake_feature_a", path_prefixes=("/feature-a",), register_fn=lambda app, module: None, ), - SimpleNamespace( + LazyFeature( name="feature-b", module_path="fake_feature_b", path_prefixes=("/feature-b",), register_fn=lambda app, module: None, ), ] - monkeypatch.setitem( - sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module - ) + monkeypatch.setitem(sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module) def fake_get_openapi(title, version, routes): path = routes[0].path @@ -58,30 +57,59 @@ def test_generate_snapshot_uses_shared_operation_id_reservations(monkeypatch): fake_proxy_server_module = ModuleType("litellm.proxy.proxy_server") fake_proxy_server_module.app = fake_app - fake_proxy_server_module.ensure_unique_openapi_operation_ids = ( - fake_ensure_unique_openapi_operation_ids - ) - monkeypatch.setitem( - sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module - ) + fake_proxy_server_module.ensure_unique_openapi_operation_ids = fake_ensure_unique_openapi_operation_ids + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module) monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi) fragments = _lazy_openapi_snapshot.generate_snapshot() - assert ( - fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["operationId"] - == "shared_operation_id_get" - ) - assert ( - fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["operationId"] - == "shared_operation_id_get_2" - ) - assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["tags"] == [ - "feature-a" - ] - assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["tags"] == [ - "feature-b" + assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["operationId"] == "shared_operation_id_get" + assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["operationId"] == "shared_operation_id_get_2" + assert fragments["feature-a"]["paths"]["/feature-a/items"]["get"]["tags"] == ["feature-a"] + assert fragments["feature-b"]["paths"]["/feature-b/items"]["get"]["tags"] == ["feature-b"] + + +def test_generate_snapshot_registers_transitively_imported_modules(monkeypatch): + """A feature module already in sys.modules (pulled in transitively by an + earlier feature) must still get register_fn called, else its routes never + mount and its fragment silently vanishes from the snapshot. Fragment + collection must also honor path_suffixes, not just prefixes.""" + from litellm.proxy import _lazy_openapi_snapshot + + fake_app = SimpleNamespace(title="LiteLLM test", version="0.0.0", routes=[]) + + fake_module = ModuleType("fake_transitive_feature") + monkeypatch.setitem(sys.modules, "fake_transitive_feature", fake_module) + + def register_fn(app, module): + app.routes.append(SimpleNamespace(path="/transitive/items")) + app.routes.append(SimpleNamespace(path="/v1/{param}/deep/leaf")) + + fake_lazy_features_module = ModuleType("litellm.proxy._lazy_features") + fake_lazy_features_module.LAZY_FEATURES = [ + LazyFeature( + name="transitive", + module_path="fake_transitive_feature", + path_prefixes=("/transitive",), + path_suffixes=("/deep/leaf",), + register_fn=register_fn, + ) ] + monkeypatch.setitem(sys.modules, "litellm.proxy._lazy_features", fake_lazy_features_module) + + def fake_get_openapi(title, version, routes): + return {"paths": {route.path: {"get": {"operationId": f"op{i}_get"}} for i, route in enumerate(routes)}} + + fake_proxy_server_module = ModuleType("litellm.proxy.proxy_server") + fake_proxy_server_module.app = fake_app + fake_proxy_server_module.ensure_unique_openapi_operation_ids = lambda schema, reserved_operation_ids: schema + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_proxy_server_module) + monkeypatch.setattr("fastapi.openapi.utils.get_openapi", fake_get_openapi) + + fragments = _lazy_openapi_snapshot.generate_snapshot() + + assert fragments["transitive"]["paths"]["/transitive/items"]["get"]["tags"] == ["transitive"] + assert "/v1/{param}/deep/leaf" in fragments["transitive"]["paths"] def test_normalize_operation_ids_uses_each_http_method(): 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 22f9e6bb67a..f48e1dba601 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -812,6 +812,71 @@ async def test_add_litellm_data_to_request_strips_callback_control_fields( assert control_field not in snapshot_body +@pytest.mark.asyncio +@pytest.mark.parametrize("timeout_field", ["timeout", "request_timeout", "stream_timeout"]) +async def test_add_litellm_data_to_request_marks_body_timeout_as_client_side(timeout_field): + """Router._get_timeout resolves the effective timeout from any of kwargs["timeout"], + kwargs["request_timeout"], or kwargs["stream_timeout"], all settable directly in the + request body. Without recognizing all three, a caller could force a 408 on every + deployment in a fallback chain without it being flagged as caller-controlled, cooling + down deployments other tenants rely on (see cooldown_handlers._trigger_cooldown_for_failed_deployment).""" + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hi"}], + timeout_field: 0.001, + }, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["client_side_timeout"] is True + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_ignores_forged_client_side_timeout(): + """The client_side_timeout marker itself must never be trusted verbatim from the + request body: a caller forging client_side_timeout=True without a real timeout + override could dodge cooldown protection on an actual deployment failure.""" + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + updated = await add_litellm_data_to_request( + data={ + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hi"}], + "client_side_timeout": True, + }, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert not updated.get("client_side_timeout") + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_allows_client_mock_response_with_admin_opt_in(): request_mock = MagicMock(spec=Request) @@ -2011,6 +2076,55 @@ def test_get_num_retries_from_request(): assert result == -1 +def test_get_keepalive_seconds_from_request(): + """ + Test LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request method + """ + # Header present with valid float string + headers_with_keepalive = {"x-litellm-keepalive-seconds": "15"} + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( + headers_with_keepalive + ) + assert result == 15.0 + + # Header not present + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( + {"Content-Type": "application/json"} + ) + assert result is None + + # Empty headers dictionary + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request({}) + assert result is None + + # Header present with a fractional value + result = LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( + {"x-litellm-keepalive-seconds": "1.5"} + ) + assert result == 1.5 + + # Header present with invalid value raises ValueError, matching the other + # x-litellm-* numeric header helpers (_get_timeout_from_request, etc.) + with pytest.raises(ValueError): + LiteLLMProxyRequestSetup._get_keepalive_seconds_from_request( + {"x-litellm-keepalive-seconds": "not-a-number"} + ) + + +def test_add_litellm_data_for_backend_llm_call_merges_keepalive_seconds_header(): + """ + The x-litellm-keepalive-seconds header must be merged into the data dict + that add_litellm_data_to_request later data.update()s onto the request body, + the same way x-litellm-timeout/x-litellm-num-retries already are. + """ + result = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={"x-litellm-keepalive-seconds": "20"}, + request_data={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + assert result.get("keepalive_seconds") == 20.0 + + def test_add_user_api_key_auth_to_request_metadata(): """ Test that add_user_api_key_auth_to_request_metadata properly adds user API key authentication data to request metadata @@ -2713,6 +2827,149 @@ def test_get_chain_id_from_headers_generic_vendor_session_id(): ) +def test_trace_id_from_traceparent_valid(): + from litellm.proxy.litellm_pre_call_utils import _trace_id_from_traceparent + + assert ( + _trace_id_from_traceparent("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01") + == "4bf92f3577b34da6a3ce929d0e0e4736" + ) + # Case-insensitive, normalized to lowercase + assert ( + _trace_id_from_traceparent("00-4BF92F3577B34DA6A3CE929D0E0E4736-00f067aa0ba902b7-01") + == "4bf92f3577b34da6a3ce929d0e0e4736" + ) + + +@pytest.mark.parametrize( + "traceparent", + [ + "not-a-traceparent", + "00-tooshort-00f067aa0ba902b7-01", + "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7", # missing flags segment + "00-4bf92f3577g34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", # non-hex char + "00-00000000000000000000000000000000-00f067aa0ba902b7-01", # all-zero trace-id, invalid per spec + "", + ], +) +def test_trace_id_from_traceparent_rejects_malformed(traceparent: str): + from litellm.proxy.litellm_pre_call_utils import _trace_id_from_traceparent + + assert _trace_id_from_traceparent(traceparent) is None + + +def test_session_id_from_baggage_valid(): + from litellm.proxy.litellm_pre_call_utils import _session_id_from_baggage + + assert _session_id_from_baggage("session.id=abc-123,user.id=42") == "abc-123" + assert _session_id_from_baggage("user.id=42, session.id=xyz-789") == "xyz-789" + + +@pytest.mark.parametrize( + "baggage", + [ + "user.id=42", + "", + "session.id=", + ], +) +def test_session_id_from_baggage_absent_or_empty(baggage: str): + from litellm.proxy.litellm_pre_call_utils import _session_id_from_baggage + + assert _session_id_from_baggage(baggage) is None + + +def test_add_litellm_metadata_from_request_headers_traceparent_sets_trace_id_only(): + """A bare traceparent header (no litellm-specific headers) sets litellm_trace_id + from its trace-id component and leaves litellm_session_id unset.""" + headers = {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert data["metadata"]["trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert "litellm_session_id" not in data + + +def test_add_litellm_metadata_from_request_headers_baggage_sets_session_id_only(): + """A bare baggage header (no litellm-specific headers) sets litellm_session_id + from its session.id entry and leaves litellm_trace_id unset.""" + headers = {"baggage": "session.id=baggage-session-42,user.id=7"} + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_session_id"] == "baggage-session-42" + assert data["metadata"]["session_id"] == "baggage-session-42" + assert "litellm_trace_id" not in data + + +def test_add_litellm_metadata_from_request_headers_baggage_session_id_not_logged_raw(caplog): + """The raw baggage session.id value must never reach the debug log line - + it isn't sanitized until set_session_id() runs much later in + Logging.__init__(), so logging it here would let a caller with control + characters or terminal escape sequences forge plaintext log output.""" + import logging + + poisoned = "poisoned\x1b[31mFAKE_RED_TEXT\x1b[0m" + headers = {"baggage": f"session.id={poisoned}"} + data = {"metadata": {}} + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_session_id"] == poisoned + assert not any(poisoned in record.getMessage() for record in caplog.records) + + +def test_add_litellm_metadata_from_request_headers_traceparent_and_baggage_together(): + """traceparent and baggage are resolved independently - trace_id and + session_id do not have to be the same value, unlike the chain_id path.""" + headers = { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "session.id=baggage-session-42", + } + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert data["litellm_session_id"] == "baggage-session-42" + + +def test_add_litellm_metadata_from_request_headers_explicit_trace_id_beats_traceparent(): + """x-litellm-trace-id must win over a traceparent header carrying a + different trace-id - explicit litellm headers are always highest priority.""" + headers = { + "x-litellm-trace-id": "explicit-trace-id-value", + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + } + data = {"metadata": {}} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers=headers, data=data, _metadata_variable_name="metadata" + ) + assert data["litellm_trace_id"] == "explicit-trace-id-value" + assert data["litellm_session_id"] == "explicit-trace-id-value" + + +def test_add_litellm_metadata_from_request_headers_anthropic_metadata_beats_baggage(): + """The existing Anthropic metadata.user_id session_id path must win over a + baggage session.id fallback.""" + data = { + "metadata": { + "user_id": "user_abc123_account__session_e96634a3-fa28-4083-b354-55542e2dca01", + } + } + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"baggage": "session.id=baggage-session-42"}, + data=data, + _metadata_variable_name="metadata", + ) + assert data["litellm_session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01" + assert "litellm_trace_id" not in data + + def test_get_internal_user_header_from_mapping_returns_expected_header(): mappings = [ {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 580a58885d9..87c8c180d9e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11278,3 +11278,68 @@ async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkey assert result is None assert mock_client.start_db_health_watchdog_task.await_count == 0 assert mock_client.health_check.await_count == 0 + + +async def _run_scheduled_background_jobs(): + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.utils import ProxyLogging + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_config = AsyncMock() + + 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={}, + 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, + ) + + import litellm.proxy.proxy_server as ps + + assert ps.scheduler is not None + return ps.scheduler + + +@pytest.mark.asyncio +async def test_ptu_rollup_job_registered_at_startup(monkeypatch): + """The PTU rollup cron is registered once an operator opts in; only models with PTU config accrue flat cost (asserted in test_ptu_flat_cost_rollup.py).""" + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( + PTU_ROLLUP_JOB_ID, + ) + + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + scheduler = await _run_scheduled_background_jobs() + + assert scheduler.get_job(PTU_ROLLUP_JOB_ID) is not None + + +@pytest.mark.asyncio +async def test_ptu_rollup_job_not_registered_without_opt_in(monkeypatch): + """Without LITELLM_ENABLE_PTU_COST_ATTRIBUTION the rollup never runs, so no sentinel row + is ever written. This is the gate that keeps the whole feature inert by default.""" + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( + PTU_ROLLUP_JOB_ID, + ) + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + scheduler = await _run_scheduled_background_jobs() + + assert scheduler.get_job(PTU_ROLLUP_JOB_ID) is None + assert len(scheduler.get_jobs()) > 0 diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 1075bffbeb2..8ee4e92ca9b 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -2928,3 +2928,143 @@ def test_update_mcp_semantic_filter_settings_requires_proxy_admin(monkeypatch): assert "proxy admin" in resp.json()["detail"].lower() finally: app.dependency_overrides.pop(user_api_key_auth, None) + + +class TestPtuCostAttributionUISetting: + """``enable_ptu_cost_attribution`` is derived from the environment on every GET. + + It is deliberately not an allowlisted, persisted setting: the point of gating PTU + flat cost on an env var is that an admin cannot flip it at runtime from the UI. + """ + + @staticmethod + def _mock_prisma(monkeypatch, stored=None): + from unittest.mock import AsyncMock, MagicMock + + mock_prisma = MagicMock() + mock_record = None + if stored is not None: + mock_record = MagicMock() + mock_record.ui_settings = stored + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=mock_record) + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + return mock_prisma + + def test_reported_false_when_the_env_var_is_unset(self, mock_auth, monkeypatch): + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is False + + def test_reported_true_once_the_env_var_is_set(self, mock_auth, monkeypatch): + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is True + + def test_a_persisted_true_cannot_forge_the_derived_value(self, mock_auth, monkeypatch): + """A row written before the allowlist existed must not be able to turn the feature on.""" + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + self._mock_prisma(monkeypatch, stored={"enable_ptu_cost_attribution": True}) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is False + + def test_is_not_an_allowlisted_persisted_setting(self): + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + ALLOWED_UI_SETTINGS_FIELDS, + ) + + assert "enable_ptu_cost_attribution" not in ALLOWED_UI_SETTINGS_FIELDS + + def test_the_body_get_returns_is_a_valid_patch_body(self, mock_auth, monkeypatch): + """Read-modify-write is how a client edits one setting. GET injects the derived key, + so rejecting it on presence made GET's own output an invalid PATCH body: the caller + got a 400 and silently lost the edit it actually wanted.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + mock_prisma = self._mock_prisma(monkeypatch) + + try: + round_tripped = client.get("/get/ui_settings").json()["values"] + assert "enable_ptu_cost_attribution" in round_tripped + response = client.patch("/update/ui_settings", json=round_tripped) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + assert mock_prisma.db.litellm_uisettings.upsert.called + + def test_a_co_submitted_setting_still_applies_alongside_the_derived_key(self, mock_auth, monkeypatch): + """The derived key riding along must not discard the caller's real edit.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + mock_prisma = self._mock_prisma(monkeypatch) + + try: + response = client.patch( + "/update/ui_settings", + json={"enable_ptu_cost_attribution": False, "enable_chat_ui": True}, + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + upsert_data = mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"] + persisted = json.loads(upsert_data["create"]["ui_settings"]) + assert persisted["enable_chat_ui"] is True + assert "enable_ptu_cost_attribution" not in persisted + + def test_patch_rejects_the_derived_setting(self, mock_auth, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + mock_prisma = self._mock_prisma(monkeypatch) + + try: + response = client.patch( + "/update/ui_settings", + json={"enable_ptu_cost_attribution": True}, + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 400 + assert "enable_ptu_cost_attribution" in str(response.json()["detail"]) + assert not mock_prisma.db.litellm_uisettings.upsert.called diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py index e33da672599..cf906259246 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py @@ -8,7 +8,6 @@ because they are direct dependents on the lifecycle state. from __future__ import annotations -import asyncio from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -17,7 +16,6 @@ import pytest import litellm from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.utils import ( - InternalUsageCache, ProxyLogging, ) @@ -102,9 +100,7 @@ def test_update_values_with_no_args_is_noop(proxy_logging): def test_update_values_invalid_type_for_alerting_raises(proxy_logging): - proxy_logging.slack_alerting_instance = MagicMock( - update_values=MagicMock(side_effect=TypeError("bad type")) - ) + proxy_logging.slack_alerting_instance = MagicMock(update_values=MagicMock(side_effect=TypeError("bad type"))) with pytest.raises(TypeError): proxy_logging.update_values(alerting={"not": "a list"}) # type: ignore[arg-type] @@ -190,6 +186,90 @@ def test_add_proxy_hooks_registers_callbacks(proxy_logging, monkeypatch): } +def _stub_hook_classes(): + class _PrismaFreeHook: + def __init__(self, internal_usage_cache): + self.internal_usage_cache = internal_usage_cache + + class _PrismaRequiringHook: + def __init__(self, internal_usage_cache, prisma_client): + self.internal_usage_cache = internal_usage_cache + self.prisma_client = prisma_client + + class _PrismaOnlyHook: + def __init__(self, prisma_client): + self.prisma_client = prisma_client + + return { + "cache_control_check": _PrismaFreeHook, + "needs_db_hook": _PrismaRequiringHook, + "db_only_hook": _PrismaOnlyHook, + } + + +def test_add_proxy_hooks_skips_prisma_requiring_hook_when_no_db(proxy_logging, monkeypatch): + hook_classes = _stub_hook_classes() + registered: List[Any] = [] + + from litellm.proxy import utils as utils_mod + + monkeypatch.setattr(utils_mod, "PROXY_HOOKS", list(hook_classes.keys())) + monkeypatch.setattr(utils_mod, "get_proxy_hook", hook_classes.__getitem__) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_callback", + lambda cb: registered.append(cb), + ) + + with patch("litellm.proxy.proxy_server.prisma_client", None): + proxy_logging._add_proxy_hooks(llm_router=None) + + snapshot = { + "mapping_keys": list(proxy_logging.proxy_hook_mapping.keys()), + "registered_types": [type(r).__name__ for r in registered], + "needs_db_hook_lookup": proxy_logging.get_proxy_hook("needs_db_hook"), + "db_only_hook_lookup": proxy_logging.get_proxy_hook("db_only_hook"), + } + assert snapshot == { + "mapping_keys": ["cache_control_check"], + "registered_types": ["_PrismaFreeHook"], + "needs_db_hook_lookup": None, + "db_only_hook_lookup": None, + } + + +def test_add_proxy_hooks_registers_prisma_requiring_hook_with_db(proxy_logging, monkeypatch): + hook_classes = _stub_hook_classes() + registered: List[Any] = [] + fake_prisma = MagicMock() + + from litellm.proxy import utils as utils_mod + + monkeypatch.setattr(utils_mod, "PROXY_HOOKS", list(hook_classes.keys())) + monkeypatch.setattr(utils_mod, "get_proxy_hook", hook_classes.__getitem__) + monkeypatch.setattr( + litellm.logging_callback_manager, + "add_litellm_callback", + lambda cb: registered.append(cb), + ) + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + proxy_logging._add_proxy_hooks(llm_router=None) + + snapshot = { + "mapping_keys": list(proxy_logging.proxy_hook_mapping.keys()), + "registered_count": len(registered), + "needs_db_hook_got_prisma": proxy_logging.proxy_hook_mapping["needs_db_hook"].prisma_client is fake_prisma, + "db_only_hook_got_prisma": proxy_logging.proxy_hook_mapping["db_only_hook"].prisma_client is fake_prisma, + } + assert snapshot == { + "mapping_keys": ["cache_control_check", "needs_db_hook", "db_only_hook"], + "registered_count": 3, + "needs_db_hook_got_prisma": True, + "db_only_hook_got_prisma": True, + } + + def test_add_proxy_hooks_unknown_hook_raises(proxy_logging, monkeypatch): from litellm.proxy import utils as utils_mod @@ -267,9 +347,7 @@ def test_init_litellm_callbacks_replaces_string_with_instance(proxy_logging, mon snapshot = { "replaced_first_item": litellm.callbacks[0] is sentinel_instance, "callbacks_grew_with_service": len(litellm.callbacks) >= 2, - "service_logging_appended": any( - "ServiceLogging" in type(c).__name__ for c in litellm.callbacks - ), + "service_logging_appended": any("ServiceLogging" in type(c).__name__ for c in litellm.callbacks), } assert snapshot == { "replaced_first_item": True, diff --git a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py index e7de8b54e4e..02ca64e5fb8 100644 --- a/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/test_litellm/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -16,10 +16,11 @@ import litellm from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( LiteLLM_ManagedVectorStore, ) -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.vector_store_endpoints.endpoints import ( _update_request_data_with_litellm_managed_vector_store_registry, index_create, + index_list, ) from litellm.proxy.vector_store_files_endpoints.endpoints import ( _update_request_data_with_model_routing_hint, @@ -37,7 +38,7 @@ from litellm.proxy.vector_store_endpoints.utils import ( is_allowed_to_call_vector_store_endpoint, is_allowed_to_call_vector_store_files_endpoint, ) -from litellm.types.vector_stores import IndexCreateRequest +from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse from litellm.types.utils import LlmProviders @@ -1316,6 +1317,93 @@ class TestIndexCreate: mock_prisma.db.litellm_managedvectorstoreindextable.create.assert_awaited_once() +class TestIndexList: + def _admin(self) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="sk-test", + key_name="sk-...test", + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + ) + + def _index_row(self, index_id: str, index_name: str) -> dict: + return { + "id": index_id, + "index_name": index_name, + "litellm_params": { + "vector_store_index": f"real-{index_name}", + "vector_store_name": "azure-ai-search", + }, + "index_info": None, + "created_at": datetime(2026, 1, 2, tzinfo=timezone.utc), + "created_by": "admin-user", + "updated_at": datetime(2026, 1, 2, tzinfo=timezone.utc), + "updated_by": "admin-user", + } + + @pytest.mark.asyncio + async def test_index_list_requires_admin(self): + """Index topology must never reach non-admins, not even via a DB read.""" + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstoreindextable.find_many = AsyncMock() + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with pytest.raises(HTTPException) as exc_info: + await index_list( + user_api_key_dict=UserAPIKeyAuth( + token="sk-test", + key_name="sk-...test", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + ) + + assert exc_info.value.status_code == 403 + assert "Only proxy admins can list" in exc_info.value.detail + mock_prisma.db.litellm_managedvectorstoreindextable.find_many.assert_not_awaited() + + @pytest.mark.asyncio + async def test_index_list_requires_db_connection(self): + with patch("litellm.proxy.proxy_server.prisma_client", None): + with pytest.raises(HTTPException) as exc_info: + await index_list(user_api_key_dict=self._admin()) + + assert exc_info.value.status_code == 500 + assert CommonProxyErrors.db_not_connected_error.value in exc_info.value.detail + + @pytest.mark.asyncio + async def test_index_list_returns_db_rows_newest_first(self): + """Rows round-trip into typed models and DB ordering (created_at desc) is requested.""" + rows = [ + self._index_row("idx-2", "index-b"), + self._index_row("idx-1", "index-a"), + ] + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstoreindextable.find_many = AsyncMock(return_value=rows) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + result = await index_list(user_api_key_dict=self._admin()) + + assert isinstance(result, IndexListResponse) + assert result.object == "list" + assert [index.index_name for index in result.data] == ["index-b", "index-a"] + assert result.data[0].litellm_params.vector_store_index == "real-index-b" + assert result.data[0].litellm_params.vector_store_name == "azure-ai-search" + assert result.data[1].litellm_params.vector_store_index == "real-index-a" + mock_prisma.db.litellm_managedvectorstoreindextable.find_many.assert_awaited_once_with( + order={"created_at": "desc"} + ) + + def test_get_v1_indexes_route_registered(self): + from litellm.proxy.vector_store_endpoints.endpoints import router + + routes = [ + (method, getattr(route, "path", None)) + for route in router.routes + for method in (getattr(route, "methods", None) or ()) + ] + assert ("GET", "/v1/indexes") in routes + + class TestIsAllowedToCallVectorStoreFilesEndpoint: def _mock_provider_config(self): provider_config = MagicMock() diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index 35f102bbb9d..c270a570ad9 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -3,7 +3,10 @@ from typing import Any, Dict, List, Mapping, Tuple import pytest -from litellm.repositories.unit_of_work import spend_reset_unit_of_work +from litellm.repositories.unit_of_work import ( + budget_cascade_unit_of_work, + spend_reset_unit_of_work, +) class FakeBatchTable: @@ -14,6 +17,9 @@ class FakeBatchTable: def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> None: self._calls.append((self._table_name, dict(where), dict(data))) + def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> None: + self._calls.append((f"{self._table_name}.update_many", dict(where), dict(data))) + class FakeBatch: def __init__(self): @@ -22,6 +28,11 @@ class FakeBatch: self.litellm_verificationtoken = FakeBatchTable("litellm_verificationtoken", self.calls) self.litellm_usertable = FakeBatchTable("litellm_usertable", self.calls) self.litellm_teamtable = FakeBatchTable("litellm_teamtable", self.calls) + self.litellm_budgettable = FakeBatchTable("litellm_budgettable", self.calls) + self.litellm_teammembership = FakeBatchTable("litellm_teammembership", self.calls) + self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls) + self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls) + self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls) async def commit(self) -> None: self.commit_count += 1 @@ -64,3 +75,53 @@ async def test_empty_block_still_commits_the_batch(): assert batch.commit_count == 1 assert batch.calls == [] + + +async def test_budget_cascade_dependents_and_window_advance_share_one_batch(): + batch = FakeBatch() + reset_at = datetime(2026, 8, 3, 12, 0, tzinfo=timezone.utc) + linked = {"budget_id": {"in": ["budget-1"]}} + + async with budget_cascade_unit_of_work(lambda: batch) as uow: + uow.team_memberships.queue_spend_zero(where=linked) + uow.keys.queue_spend_zero(where=linked) + uow.organizations.queue_spend_zero(where=linked) + uow.tags.queue_spend_zero(where=linked) + uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}}) + uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at) + assert batch.commit_count == 0 + + assert batch.commit_count == 1 + assert batch.calls == [ + ("litellm_teammembership.update_many", linked, {"spend": 0}), + ("litellm_verificationtoken.update_many", linked, {"spend": 0}), + ("litellm_organizationtable.update_many", linked, {"spend": 0}), + ("litellm_tagtable.update_many", linked, {"spend": 0}), + ("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}), + ("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}), + ] + + +async def test_budget_window_advance_tolerates_a_tier_deleted_mid_chunk(): + """A tier deleted between the read and the commit must not abort the batch: + ``update`` raises P2025 on a missing row and takes every other write in the + chunk down with it, while ``update_many`` just matches nothing.""" + batch = FakeBatch() + + async with budget_cascade_unit_of_work(lambda: batch) as uow: + uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=datetime.now(timezone.utc)) + + assert [call[0] for call in batch.calls] == ["litellm_budgettable.update_many"] + + +async def test_budget_cascade_raising_inside_block_skips_commit(): + """A failure part-way through must leave budget_reset_at where it was, so + the tier is still due on the next tick.""" + batch = FakeBatch() + + with pytest.raises(RuntimeError, match="boom"): + async with budget_cascade_unit_of_work(lambda: batch) as uow: + uow.team_memberships.queue_spend_zero(where={"budget_id": {"in": ["budget-1"]}}) + raise RuntimeError("boom") + + assert batch.commit_count == 0 diff --git a/tests/test_litellm/router_utils/test_cooldown_cache.py b/tests/test_litellm/router_utils/test_cooldown_cache.py index 52fe151eff4..a48402684b4 100644 --- a/tests/test_litellm/router_utils/test_cooldown_cache.py +++ b/tests/test_litellm/router_utils/test_cooldown_cache.py @@ -4,6 +4,7 @@ Unit tests for CooldownCache exception masking functionality import os import sys +import time from unittest.mock import MagicMock import pytest @@ -94,9 +95,7 @@ class TestCooldownCacheExceptionMasking: assert "magical kingdom" not in masked_exception # Should preserve the error type information at the beginning (first 50 chars) - assert masked_exception.startswith( - "litellm.proxy.proxy_server._handle_llm_api_excepti" - ) + assert masked_exception.startswith("litellm.proxy.proxy_server._handle_llm_api_excepti") def test_exception_with_api_keys_masked(self, cooldown_cache): """Test that API keys in exceptions are properly masked""" @@ -119,9 +118,7 @@ class TestCooldownCacheExceptionMasking: masked_exception = cooldown_data["exception_received"] # Should mask the sensitive content while preserving structure - assert masked_exception.startswith( - "Authentication failed with api_key=sk-12345678" - ) + assert masked_exception.startswith("Authentication failed with api_key=sk-12345678") assert "*" in masked_exception assert len(masked_exception) == len(exception_with_key) @@ -179,9 +176,7 @@ class TestCooldownCacheExceptionMasking: # Should successfully convert exception to string assert isinstance(cooldown_data["exception_received"], str) - assert ( - str(exc) == cooldown_data["exception_received"] - ) # Short exceptions not masked + assert str(exc) == cooldown_data["exception_received"] # Short exceptions not masked def test_masking_preserves_error_debugging_info(self, cooldown_cache): """Test that masking preserves essential debugging information""" @@ -208,9 +203,7 @@ class TestCooldownCacheExceptionMasking: masked_exception = cooldown_data["exception_received"] # Should preserve error type and initial debugging info (first 50 chars) - assert masked_exception.startswith( - "RateLimitError: Rate limit exceeded for model gpt-" - ) + assert masked_exception.startswith("RateLimitError: Rate limit exceeded for model gpt-") # Should mask the prompt content assert "Write a comprehensive analysis" not in masked_exception @@ -255,3 +248,190 @@ class TestCooldownCacheExceptionMasking: # Should show first 50 characters, then all asterisks expected = "A" * 50 + "*" * 50 assert masked == expected + + +class TestCooldownCacheTTLCorrection: + def _make_cooldown_cache(self) -> CooldownCache: + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory) + return CooldownCache(cache=dual_cache, default_cooldown_time=60.0) + + def test_expired_entry_evicted_and_not_returned(self): + """ + An entry with timestamp+cooldown_time in the past must be evicted from + in-memory cache and excluded from the active cooldown list. + """ + cc = self._make_cooldown_cache() + model_id = "expired-deployment" + key = CooldownCache.get_cooldown_cache_key(model_id) + + expired_value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time() - 120.0, + "cooldown_time": 60.0, + } + cc.cache.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" + + def test_active_entry_is_returned(self): + """ + An entry whose cooldown window has not elapsed must appear in the active list. + """ + cc = self._make_cooldown_cache() + model_id = "active-deployment" + key = CooldownCache.get_cooldown_cache_key(model_id) + + active_value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time(), + "cooldown_time": 60.0, + } + cc.cache.in_memory_cache.set_cache(key, active_value, ttl=60) + + active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None) + + assert len(active) == 1 + assert active[0][0] == model_id + + def test_ttl_corrected_when_in_memory_expiry_far_exceeds_remaining(self): + """ + When DualCache backfills from Redis using the default 600s TTL, the in-memory + TTL must be corrected to min(remaining, 60) seconds. + """ + cc = self._make_cooldown_cache() + model_id = "backfilled-deployment" + key = CooldownCache.get_cooldown_cache_key(model_id) + + remaining = 30.0 + value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time() - (60.0 - remaining), + "cooldown_time": 60.0, + } + cc.cache.in_memory_cache.set_cache(key, value, ttl=600) + + before_expiry = cc.cache.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) + assert after_expiry is not None + corrected_remaining = after_expiry - time.time() + assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s" + assert corrected_remaining > 0, "Corrected TTL must be positive (cooldown still active)" + + @pytest.mark.asyncio + async def test_async_expired_entry_evicted(self): + """ + Async path must also evict expired entries. + """ + cc = self._make_cooldown_cache() + model_id = "async-expired" + key = CooldownCache.get_cooldown_cache_key(model_id) + + expired_value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time() - 120.0, + "cooldown_time": 60.0, + } + cc.cache.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 + + @pytest.mark.asyncio + async def test_async_active_entry_is_returned(self): + """ + Async counterpart of test_active_entry_is_returned: an entry whose cooldown + window has not elapsed must appear in the async active list too. + """ + cc = self._make_cooldown_cache() + model_id = "async-active-deployment" + key = CooldownCache.get_cooldown_cache_key(model_id) + + active_value: CooldownCacheValue = { + "exception_received": "Rate limit", + "status_code": "429", + "timestamp": time.time(), + "cooldown_time": 60.0, + } + cc.cache.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) + + assert len(active) == 1 + assert active[0][0] == model_id + + +class TestCorrectedActiveCooldown: + def _make_cooldown_cache(self) -> CooldownCache: + in_memory = InMemoryCache() + dual_cache = DualCache(in_memory_cache=in_memory) + return CooldownCache(cache=dual_cache, default_cooldown_time=60.0) + + def _entry(self, timestamp: float, cooldown_time: float) -> CooldownCacheValue: + return CooldownCacheValue( + exception_received="Rate limit", + status_code="429", + timestamp=timestamp, + cooldown_time=cooldown_time, + ) + + def test_expired_entry_returns_none_and_evicts(self): + 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) + + 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 + + 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) + + result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time()) + + assert result is not None + assert result["status_code"] == "429" + + def test_inflated_ttl_is_corrected(self): + cc = self._make_cooldown_cache() + 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) + + 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) + assert corrected_expiry is not None + assert corrected_expiry - time.time() <= 60.0 + + def test_normal_ttl_not_modified(self): + 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._corrected_active_cooldown(key, dict(entry), current_time=time.time()) + + after_expiry = cc.cache.in_memory_cache.ttl_dict.get(key) + assert after_expiry == original_expiry diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py new file mode 100644 index 00000000000..2139521d2a8 --- /dev/null +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -0,0 +1,298 @@ +from unittest.mock import MagicMock, patch + +import litellm +from litellm.router_utils.cooldown_handlers import ( + _get_deployment_cooldown_policy, + _resolve_allowed_fails_from_policy, + _should_cooldown_based_on_deployment_policy, + should_cooldown_based_on_allowed_fails_policy, +) + + +class TestGetDeploymentCooldownPolicy: + def _make_router(self, deployment_id: str, model_info: dict | None = None): + router = MagicMock() + if model_info is None: + router.get_model_info.return_value = None + else: + router.get_model_info.return_value = {"model_info": model_info} + return router + + def test_deployment_not_found_returns_none_none(self): + router = self._make_router("dep-1") + policy, allowed = _get_deployment_cooldown_policy(router, "dep-1") + assert policy is None + assert allowed is None + + def test_no_model_info_returns_none_none(self): + router = MagicMock() + router.get_model_info.return_value = {"model_info": {}} + policy, allowed = _get_deployment_cooldown_policy(router, "dep-1") + assert policy is None + assert allowed is None + + def test_returns_policy_dict_and_allowed_fails(self): + router = self._make_router( + "dep-1", + {"allowed_fails_policy": {"RateLimitErrorAllowedFails": 2}, "allowed_fails": 3}, + ) + policy, allowed = _get_deployment_cooldown_policy(router, "dep-1") + assert policy == {"RateLimitErrorAllowedFails": 2} + assert allowed == 3 + + def test_non_dict_policy_treated_as_none(self): + router = self._make_router("dep-1", {"allowed_fails_policy": "invalid", "allowed_fails": 5}) + policy, allowed = _get_deployment_cooldown_policy(router, "dep-1") + assert policy is None + assert allowed == 5 + + def test_allowed_fails_only(self): + router = self._make_router("dep-1", {"allowed_fails": 1}) + policy, allowed = _get_deployment_cooldown_policy(router, "dep-1") + assert policy is None + assert allowed == 1 + + +class TestResolveAllowedFailsFromPolicy: + def test_none_policy_returns_none(self): + exc = litellm.RateLimitError("429", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(None, exc) is None + + def test_matching_rate_limit_error(self): + policy = {"RateLimitErrorAllowedFails": 3} + exc = litellm.RateLimitError("429", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(policy, exc) == 3 + + def test_matching_internal_server_error(self): + policy = {"InternalServerErrorAllowedFails": 5} + exc = litellm.InternalServerError("500", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(policy, exc) == 5 + + def test_matching_service_unavailable_error(self): + policy = {"ServiceUnavailableErrorAllowedFails": 4} + exc = litellm.ServiceUnavailableError("503", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(policy, exc) == 4 + + def test_matching_bad_gateway_error(self): + policy = {"BadGatewayErrorAllowedFails": 2} + exc = litellm.BadGatewayError("502", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(policy, exc) == 2 + + def test_matching_not_found_error(self): + policy = {"NotFoundErrorAllowedFails": 1} + exc = litellm.NotFoundError("404", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(policy, exc) == 1 + + def test_unmatched_exception_returns_none(self): + policy = {"RateLimitErrorAllowedFails": 3} + exc = litellm.InternalServerError("500", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(policy, exc) is None + + def test_field_absent_from_policy_returns_none(self): + policy: dict[str, int] = {} + exc = litellm.InternalServerError("500", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(policy, exc) is None + + def test_content_policy_violation_not_shadowed_by_bad_request_error(self): + """ContentPolicyViolationError subclasses BadRequestError, so if + BadRequestError were checked first, this would incorrectly resolve to + BadRequestErrorAllowedFails (10) instead of + ContentPolicyViolationErrorAllowedFails (2).""" + policy = {"BadRequestErrorAllowedFails": 10, "ContentPolicyViolationErrorAllowedFails": 2} + exc = litellm.ContentPolicyViolationError("flagged content", "openai", "gpt-4") + assert _resolve_allowed_fails_from_policy(policy, exc) == 2 + + +class TestShouldCooldownBasedOnDeploymentPolicy: + def _make_router(self, model_info: dict | None = None): + router = MagicMock() + if model_info is None: + router.get_model_info.return_value = None + else: + router.get_model_info.return_value = model_info + return router + + def test_policy_match_uses_exception_type_as_cache_key_suffix(self): + policy = {"RateLimitErrorAllowedFails": 0} + exc = litellm.RateLimitError("429", "openai", "gpt-4") + router = self._make_router({"litellm_params": {}, "model_info": {}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = True + result = _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, policy, None, is_single_deployment_model_group=False + ) + + assert result is True + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["allowed_fails_override"] == 0 + assert call_kwargs["cache_key_suffix"] == "RateLimitError" + + def test_no_policy_match_uses_dep_allowed_fails_and_generic_suffix(self): + policy: dict[str, int] = {} + exc = litellm.InternalServerError("500", "openai", "gpt-4") + router = self._make_router({"litellm_params": {}, "model_info": {}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = False + result = _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, policy, dep_allowed_fails=3, is_single_deployment_model_group=False + ) + + assert result is False + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["allowed_fails_override"] == 3 + assert call_kwargs["cache_key_suffix"] == "generic" + + def test_dep_allowed_fails_on_single_deployment_group_does_not_cooldown(self): + """A generic, deployment-wide allowed_fails predates the per-exception-type + policy and is a less deliberate opt-in, so on a single-deployment model group + it must not silently disable the "avoid cooldowns on single deployment model + groups" safety net.""" + exc = litellm.InternalServerError("500", "openai", "gpt-4") + router = self._make_router({"litellm_params": {}, "model_info": {}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + result = _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, None, dep_allowed_fails=3, is_single_deployment_model_group=True + ) + + assert result is False + mock_sc.assert_not_called() + + def test_named_policy_on_single_deployment_group_still_cools_down(self): + """Unlike a generic allowed_fails, an explicit per-exception-type policy entry + is a deliberate opt-in and must still apply on a single-deployment group.""" + policy = {"RateLimitErrorAllowedFails": 0} + exc = litellm.RateLimitError("429", "openai", "gpt-4") + router = self._make_router({"litellm_params": {}, "model_info": {}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = True + result = _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, policy, None, is_single_deployment_model_group=True + ) + + assert result is True + mock_sc.assert_called_once() + + def test_no_policy_and_no_dep_allowed_fails_defers_to_router_level(self): + """When neither a deployment policy nor a deployment-wide allowed_fails covers + this exception, defer to router-level behavior instead of forcing an + immediate cooldown (allowed_fails_override=0 would trip on the first failure + of any exception type the deployment's config doesn't mention).""" + exc = litellm.InternalServerError("500", "openai", "gpt-4") + router = self._make_router({"litellm_params": {}, "model_info": {}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = True + _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, None, None, is_single_deployment_model_group=False + ) + + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["allowed_fails_override"] is None + assert call_kwargs["cache_key_suffix"] is None + + def test_partial_policy_without_dep_allowed_fails_defers_for_uncovered_exception(self): + """A deployment that only sets RateLimitErrorAllowedFails must not force a + zero-fail threshold on an unrelated TimeoutError; it should defer to + router-level behavior for exception types its policy doesn't mention.""" + policy = {"RateLimitErrorAllowedFails": 0} + exc = litellm.Timeout("timed out", "openai", "gpt-4") + router = self._make_router({"litellm_params": {}, "model_info": {}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = False + _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, policy, dep_allowed_fails=None, is_single_deployment_model_group=False + ) + + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["allowed_fails_override"] is None + assert call_kwargs["cache_key_suffix"] is None + + def test_cooldown_time_from_model_info_passed_through(self): + exc = litellm.RateLimitError("429", "openai", "gpt-4") + router = self._make_router({"litellm_params": {}, "model_info": {"cooldown_time": 120.0}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = True + _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, None, None, is_single_deployment_model_group=False + ) + + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["cooldown_time_override"] == 120.0 + + def test_cooldown_time_from_litellm_params_used_as_fallback(self): + """cooldown_time has pre-existing litellm_params support on the primary + failure path, so it must still be honored here when model_info doesn't + set it.""" + exc = litellm.RateLimitError("429", "openai", "gpt-4") + router = self._make_router({"litellm_params": {"cooldown_time": 120.0}, "model_info": {}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = True + _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, None, None, is_single_deployment_model_group=False + ) + + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["cooldown_time_override"] == 120.0 + + def test_cooldown_time_from_model_info_takes_priority_over_litellm_params(self): + exc = litellm.RateLimitError("429", "openai", "gpt-4") + router = self._make_router({"litellm_params": {"cooldown_time": 120.0}, "model_info": {"cooldown_time": 15.0}}) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = True + _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, None, None, is_single_deployment_model_group=False + ) + + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["cooldown_time_override"] == 15.0 + + def test_model_info_none_passes_none_cooldown_time(self): + exc = litellm.RateLimitError("429", "openai", "gpt-4") + router = self._make_router(None) + + with patch("litellm.router_utils.cooldown_handlers.should_cooldown_based_on_allowed_fails_policy") as mock_sc: + mock_sc.return_value = False + _should_cooldown_based_on_deployment_policy( + router, "dep-1", exc, None, None, is_single_deployment_model_group=False + ) + + call_kwargs = mock_sc.call_args[1] + assert call_kwargs["cooldown_time_override"] is None + + +class TestShouldCooldownBasedOnAllowedFailsPolicy: + def _make_router(self, cooldown_time: float = 60.0) -> 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 + 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) + exc = litellm.RateLimitError("429", "openai", "gpt-4") + + should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="dep-1", + original_exception=exc, + allowed_fails_override=5, + 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, ( + "cooldown_time_override=0 should be used as TTL, not the router-level 60.0" + ) diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index e2348d28701..68395737469 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -1,9 +1,14 @@ import json +from unittest.mock import MagicMock, patch +import httpx import pytest +import litellm +from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure from litellm.router_utils.fallback_event_handlers import ( AttemptedFallbackTargets, + _trigger_cooldown_for_failed_deployment, fallback_attempt_key, get_fallback_model_group, run_async_fallback, @@ -144,6 +149,147 @@ async def test_run_async_fallback_skips_original_model_group(): assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 +class AttemptRecordingRouter: + def __init__(self): + self.attempted_model_groups = [] + self.received_kwargs = None + + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + self.attempted_model_groups.append(kwargs.get("model")) + self.received_kwargs = kwargs + return StreamingWrapper() + + +async def _acreate_batch(*args, **kwargs): + raise AssertionError("only used for its __name__") + + +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_uploaded_file_requests_in_their_model_group(): + """An input_file_id only exists under the credentials of the group it was uploaded + to, so a cross-group fallback can only fail with the wrong provider's error.""" + router = AttemptRecordingRouter() + owning_provider_error = RuntimeError("openai connection error") + + with pytest.raises(RuntimeError, match="openai connection error"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=owning_provider_error, + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + input_file_id="file-owned-by-openai", + original_function=_acreate_batch, + ) + + assert router.attempted_model_groups == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_fine_tuning_requests_in_their_model_group(): + router = AttemptRecordingRouter() + + with pytest.raises(RuntimeError, match="openai connection error"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + training_file="file-owned-by-openai", + ) + + assert router.attempted_model_groups == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_allows_same_model_group_retry_for_uploaded_file_requests(): + """Order-based fallbacks stay inside the owning group, so they must still run.""" + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "openai-group", "_target_order": 2}], + original_model_group="openai-group", + original_exception=RuntimeError("first deployment failed"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + input_file_id="file-owned-by-openai", + original_function=_acreate_batch, + ) + + assert router.attempted_model_groups == ["openai-group"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_still_crosses_model_groups_without_an_uploaded_file(): + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + ) + + assert router.attempted_model_groups == ["azure-group"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_handles_explicitly_none_metadata(): + """/v1/batches always sets `metadata`, and sets it to None when the caller sent + none, so setdefault() on it hands back None instead of a dict.""" + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + metadata=None, + ) + + assert router.received_kwargs["metadata"] == {"model_group": "azure-group"} + + +@pytest.mark.asyncio +async def test_run_async_fallback_records_batch_model_group_outside_provider_metadata(): + """`metadata` on a batch request is forwarded to the provider and stored on the + batch, so the router's own model_group belongs in litellm_metadata.""" + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "openai-group", "_target_order": 2}], + original_model_group="openai-group", + original_exception=RuntimeError("first deployment failed"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + input_file_id="file-owned-by-openai", + metadata={"caller": "nightly-job"}, + litellm_metadata={"model_group": "openai-group"}, + original_function=_acreate_batch, + ) + + assert router.received_kwargs["metadata"] == {"caller": "nightly-job"} + assert router.received_kwargs["litellm_metadata"]["model_group"] == "openai-group" + + class RecordingFailRouter: def __init__(self): self.attempted_models = [] @@ -351,9 +497,304 @@ def test_get_fallback_model_group_does_not_mutate_fallbacks(): fallbacks list, which is the live router config shared across requests.""" fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}, "gpt-4o-mini"] - fallback_model_group, _ = get_fallback_model_group( - fallbacks=fallbacks, model_group="unmatched-model" - ) + fallback_model_group, _ = get_fallback_model_group(fallbacks=fallbacks, model_group="unmatched-model") assert fallback_model_group == ["gpt-4o-mini"] assert fallbacks == [{"gpt-3.5-turbo": ["claude-3-haiku"]}, "gpt-4o-mini"] + + +class TestTriggerCooldownForFailedDeployment: + def test_calls_set_cooldown_deployments_with_stamped_deployment_id(self): + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + mock_set_cooldown.assert_called_once() + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["deployment"] == "fallback-deployment" + assert call_kwargs["original_exception"] is exc + + def test_does_not_trust_caller_supplied_metadata_bucket(self): + """A metadata bucket can't reliably be told apart from a caller-supplied + one without knowing this call's function_name, so a client with + permission to set metadata must not be able to get an arbitrary + deployment cooled down by forging a deployment_model_name marker.""" + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + kwargs = { + "metadata": { + "model_info": {"id": "attacker-chosen-deployment"}, + "deployment_model_name": "gpt-4", + } + } + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs=kwargs, exception=exc) + + mock_set_cooldown.assert_not_called() + + def test_increments_failure_counter_before_cooldown_check(self): + """The fallback path must feed the same per-minute failure counter the + primary path uses, or repeated fallback failures never accumulate + toward the default percent-fail-rate cooldown threshold.""" + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with ( + patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown, + patch( + "litellm.router_utils.fallback_event_handlers.increment_deployment_failures_for_current_minute" + ) as mock_increment, + ): + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + mock_increment.assert_called_once_with( + litellm_router_instance=mock_router, deployment_id="fallback-deployment" + ) + mock_set_cooldown.assert_called_once() + + def test_no_op_when_deployment_id_missing(self): + mock_router = MagicMock() + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, kwargs={}, exception=RuntimeError("no metadata") + ) + + mock_set_cooldown.assert_not_called() + + def test_skipped_for_advisor_orchestration_failure(self): + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + mark_advisor_orchestration_failure(exc) + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + mock_set_cooldown.assert_not_called() + + def test_uses_deployment_litellm_params_cooldown_time_override(self): + mock_router = MagicMock() + mock_router.cooldown_time = 300.0 + mock_router.get_model_info.return_value = {"litellm_params": {"cooldown_time": 30.0}} + + exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["time_to_cooldown"] == 30.0 + + def test_uses_response_header_when_no_deployment_config(self): + """Precedence must match Router.deployment_callback_on_failure's primary + path: deployment config, then the response's Retry-After header, then the + router default.""" + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = {"litellm_params": {}} + + exc = RuntimeError("upstream error") + exc.failed_deployment_id = "fallback-deployment" + exc.litellm_response_headers = httpx.Headers({"retry-after": "45"}) + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + call_kwargs = mock_set_cooldown.call_args[1] + assert call_kwargs["time_to_cooldown"] == 45 + + def test_silently_catches_exceptions(self): + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = RuntimeError("upstream error") + exc.failed_deployment_id = "fallback-deployment" + + with patch( + "litellm.router_utils.fallback_event_handlers._set_cooldown_deployments", + side_effect=RuntimeError("cooldown error"), + ): + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + def test_skips_request_scoped_404_on_generic_api_call(self): + """A generic API call (files/batches/threads/rerank/...) forwards a caller-supplied + resource id, so a 404 there means "that id doesn't exist", not "this deployment is + unhealthy". Without this guard, a single bad id would 404 every deployment in the + fallback chain and cool all of them down from one request.""" + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.NotFoundError("not found", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with ( + patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown, + patch( + "litellm.router_utils.fallback_event_handlers.increment_deployment_failures_for_current_minute" + ) as mock_increment, + ): + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs={"original_generic_function": MagicMock()}, + exception=exc, + ) + + mock_set_cooldown.assert_not_called() + mock_increment.assert_not_called() + + def test_still_cools_down_404_outside_generic_api_call(self): + """The request-scoped-404 guard is scoped to generic API calls only: a 404 on a + regular completion fallback (no original_generic_function in kwargs) must still + cool down the deployment as before.""" + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.NotFoundError("not found", "openai", "gpt-4") + exc.failed_deployment_id = "fallback-deployment" + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + mock_set_cooldown.assert_called_once() + + def test_skips_client_side_timeout_408(self): + """The proxy's x-litellm-timeout header lets a caller set an arbitrarily short + timeout, which litellm.Timeout reports as status 408 regardless of the + deployment's actual health. Without this guard, a caller could force a 408 on + every deployment in the fallback chain from a single request.""" + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.Timeout(message="timeout", model="gpt-4", llm_provider="openai") + exc.failed_deployment_id = "fallback-deployment" + + with ( + patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown, + patch( + "litellm.router_utils.fallback_event_handlers.increment_deployment_failures_for_current_minute" + ) as mock_increment, + ): + _trigger_cooldown_for_failed_deployment( + litellm_router=mock_router, + kwargs={"client_side_timeout": True}, + exception=exc, + ) + + mock_set_cooldown.assert_not_called() + mock_increment.assert_not_called() + + def test_still_cools_down_408_without_client_side_timeout_flag(self): + """The client-side-timeout guard is scoped to caller-supplied timeouts only: a + 408 that did not come from x-litellm-timeout (no client_side_timeout in kwargs) + must still cool down the deployment as before.""" + mock_router = MagicMock() + mock_router.cooldown_time = 60.0 + mock_router.get_model_info.return_value = None + + exc = litellm.Timeout(message="timeout", model="gpt-4", llm_provider="openai") + exc.failed_deployment_id = "fallback-deployment" + + with patch("litellm.router_utils.fallback_event_handlers._set_cooldown_deployments") as mock_set_cooldown: + _trigger_cooldown_for_failed_deployment(litellm_router=mock_router, kwargs={}, exception=exc) + + mock_set_cooldown.assert_called_once() + + +class TestRunAsyncFallbackTriggersCooldown: + class RouterWithLoggingKwarg: + def __init__(self): + self.cooldown_time = 60.0 + + def log_retry(self, kwargs, e): + return kwargs + + def get_model_info(self, id): + return None + + async def async_function_with_fallbacks(self, *args, **kwargs): + raise RuntimeError("fallback model also failed") + + def _logging_obj(self, has_logged_async_failure: bool) -> MagicMock: + logging_obj = MagicMock() + logging_obj.model_call_details = {"has_logged_async_failure": has_logged_async_failure} + return logging_obj + + @pytest.mark.asyncio + async def test_triggers_cooldown_when_has_logged_async_failure_is_true(self): + with patch( + "litellm.router_utils.fallback_event_handlers._trigger_cooldown_for_failed_deployment" + ) as mock_trigger: + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=self.RouterWithLoggingKwarg(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original request failed"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=self._logging_obj(has_logged_async_failure=True), + ) + + mock_trigger.assert_called_once() + + @pytest.mark.asyncio + async def test_does_not_trigger_cooldown_when_has_logged_async_failure_is_false(self): + """This is the exact dead-code scenario the bug fix addresses: before it, + the normal failure callback runs for the first attempt in a fallback chain + (has_logged_async_failure is still False at that point), so no explicit + trigger is needed there.""" + with patch( + "litellm.router_utils.fallback_event_handlers._trigger_cooldown_for_failed_deployment" + ) as mock_trigger: + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=self.RouterWithLoggingKwarg(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original request failed"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=self._logging_obj(has_logged_async_failure=False), + ) + + mock_trigger.assert_not_called() + + @pytest.mark.asyncio + async def test_does_not_trigger_cooldown_when_no_logging_obj_present(self): + with patch( + "litellm.router_utils.fallback_event_handlers._trigger_cooldown_for_failed_deployment" + ) as mock_trigger: + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=self.RouterWithLoggingKwarg(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original request failed"), + max_fallbacks=3, + fallback_depth=0, + ) + + mock_trigger.assert_not_called() diff --git a/tests/test_litellm/router_utils/test_router_utils_common_utils.py b/tests/test_litellm/router_utils/test_router_utils_common_utils.py index 0d063ad14f5..30f658d7ea2 100644 --- a/tests/test_litellm/router_utils/test_router_utils_common_utils.py +++ b/tests/test_litellm/router_utils/test_router_utils_common_utils.py @@ -1,3 +1,4 @@ +import logging from typing import Dict, List, Optional, Union from unittest.mock import Mock @@ -13,6 +14,8 @@ from litellm.router_utils.common_utils import ( filter_web_search_deployments, resolve_model_group_alias, truncate_fallback_error_detail, + PROVIDER_SCOPED_CREDENTIAL_PARAMS, + warn_on_provider_credential_mismatch, ) @@ -584,3 +587,172 @@ class TestTruncateFallbackErrorDetail: to stay small enough that a walk over many model groups cannot compound it into an output volume that starves the process.""" assert len(truncate_fallback_error_detail("x" * 1_000_000)) < 3_000 + + +class TestWarnOnProviderCredentialMismatch: + """A deployment that carries one provider's credentials while resolving to + another is silently broken: litellm ignores the credentials and sends the + request to the resolved provider, which 401s. The classic shape is a bedrock + model group where one entry lost its route prefix, which fails only on the + requests the router happens to send to that entry.""" + + def test_warns_when_aws_params_sit_on_an_anthropic_model(self): + warning = warn_on_provider_credential_mismatch( + model_name="claude-sonnet-5", + litellm_params={"model": "claude-sonnet-5", "aws_region_name": "eu-central-1"}, + ) + + assert warning is not None + assert "aws_region_name" in warning + assert "anthropic" in warning + assert "bedrock/claude-sonnet-5" in warning + + def test_silent_when_the_prefix_is_present(self): + assert ( + warn_on_provider_credential_mismatch( + model_name="claude-sonnet-5", + litellm_params={ + "model": "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_region_name": "eu-central-1", + }, + ) + is None + ) + + def test_silent_when_custom_llm_provider_supplies_the_route(self): + """An operator may name the provider explicitly instead of prefixing the + model; that is consistent and must not warn.""" + assert ( + warn_on_provider_credential_mismatch( + model_name="claude-sonnet-5", + litellm_params={ + "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", + "custom_llm_provider": "bedrock", + "aws_region_name": "eu-central-1", + }, + ) + is None + ) + + def test_silent_when_no_provider_scoped_credentials_are_set(self): + assert ( + warn_on_provider_credential_mismatch( + model_name="gpt-5.5", litellm_params={"model": "gpt-5.5"} + ) + is None + ) + + def test_vertex_params_name_vertex_not_bedrock(self): + """The hint must follow the params that were actually set, otherwise it + sends the operator to the wrong prefix.""" + warning = warn_on_provider_credential_mismatch( + model_name="claude-on-vertex", + litellm_params={"model": "claude-sonnet-5", "vertex_project": "my-project"}, + ) + + assert warning is not None + assert "vertex_ai/claude-sonnet-5" in warning + assert "bedrock" not in warning + + def test_silent_for_a_model_litellm_cannot_classify(self): + """An unresolvable model must not warn and must not raise: this runs on + the router startup path, so a wrong guess would spam every boot.""" + assert ( + warn_on_provider_credential_mismatch( + model_name="mystery", + litellm_params={"model": "not-a-real-provider-model-xyz", "aws_region_name": "us-east-1"}, + ) + is None + ) + + def test_router_warns_for_a_config_shaped_model_list(self, caplog): + """The whole point is that this fires where operators declare models, so + drive Router rather than the helper.""" + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + Router( + model_list=[ + { + "model_name": "claude-sonnet-5", + "litellm_params": { + "model": "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_region_name": "us-east-1", + }, + }, + { + "model_name": "claude-sonnet-5", + "litellm_params": { + "model": "claude-sonnet-5", + "aws_region_name": "us-east-1", + }, + }, + ] + ) + + mismatch_warnings = [r for r in caplog.records if "resolves to provider" in r.getMessage()] + assert len(mismatch_warnings) == 1, ( + "exactly the prefix-less deployment should warn; " + f"got {[r.getMessage() for r in mismatch_warnings]}" + ) + assert "aws_region_name" in mismatch_warnings[0].getMessage() + + @pytest.mark.parametrize( + "model", + [ + "bedrock/mantle/anthropic.claude-sonnet-4-5-20250929-v1:0", + "bedrock/converse/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "sagemaker/my-endpoint", + ], + ) + def test_silent_for_every_aws_family_route(self, model): + """The AWS family is wider than 'bedrock': mantle, sagemaker and the + sagemaker variants all read aws_* legitimately. Warning on any of them + would tell an operator to 'fix' a working deployment, so the provider + set is derived from LlmProviders rather than hand-listed.""" + assert ( + warn_on_provider_credential_mismatch( + model_name="aws-deployment", + litellm_params={"model": model, "aws_region_name": "us-east-1"}, + ) + is None + ) + + def test_every_aws_family_provider_is_covered(self): + """Pins the derivation itself: a newly added bedrock_*/sagemaker_* provider + must join the set automatically, or it starts drawing false warnings.""" + from litellm.types.utils import LlmProviders + + aws_family = {p.value for p in LlmProviders if p.value.startswith(("bedrock", "sagemaker"))} + assert aws_family <= PROVIDER_SCOPED_CREDENTIAL_PARAMS["aws_region_name"] + assert {"bedrock", "bedrock_mantle", "sagemaker", "sagemaker_chat", "sagemaker_nova"} <= aws_family + + def test_silent_when_credentials_come_from_a_named_credential(self): + """Named credentials resolve after registration, so the params are absent + here. Warning on that absence would fire on every such deployment.""" + assert ( + warn_on_provider_credential_mismatch( + model_name="claude-sonnet-5", + litellm_params={ + "model": "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "litellm_credential_name": "my-aws-creds", + }, + ) + is None + ) + + @pytest.mark.parametrize("provider", ["bedrock_mantle", "sagemaker_nova"]) + def test_silent_for_aws_providers_named_explicitly(self, provider): + """The false-positive shape: an operator names a less common AWS provider + directly, so the model string carries no route prefix to key off. A + hand-listed provider set misses these and tells them to 'fix' a working + deployment by prefixing it with bedrock/.""" + assert ( + warn_on_provider_credential_mismatch( + model_name="aws-deployment", + litellm_params={ + "model": "anthropic.claude-sonnet-4-5-20250929-v1:0", + "custom_llm_provider": provider, + "aws_region_name": "us-east-1", + }, + ) + is None + ) diff --git a/tests/test_litellm/test_check_type_discipline.py b/tests/test_litellm/test_check_type_discipline.py index 4b8533df604..25131088e9a 100644 --- a/tests/test_litellm/test_check_type_discipline.py +++ b/tests/test_litellm/test_check_type_discipline.py @@ -117,6 +117,17 @@ def test_typing_alias_and_forward_ref_annotations_are_flagged(tmp_path): assert "LIT001" in _codes(tmp_path, 'x: "dict[str, int]"\n') +def test_literal_string_args_are_values_not_forward_refs(tmp_path): + assert "LIT001" not in _codes(tmp_path, 'from typing import Literal\nx: Literal["list"] = "list"\n') + assert "LIT001" not in _codes( + tmp_path, + 'from typing import Literal\ndef f(op: Literal["create", "list"] = "create") -> None:\n return None\n', + ) + assert "LIT001" not in _codes(tmp_path, 'import typing\nx: typing.Literal["dict"] = "dict"\n') + assert "LIT001" in _codes(tmp_path, 'from typing import Literal\nx: dict[str, Literal["a"]]\n') + assert "LIT001" in _codes(tmp_path, "x: \"Literal['x'] | list[int]\"\n") + + def test_readonly_annotations_are_clean(tmp_path): for ann in ("Mapping[str, int]", "Sequence[int]", "tuple[int, ...]", "frozenset[int]"): assert "LIT001" not in _codes(tmp_path, f"from typing import Mapping, Sequence\nx: {ann}\n") diff --git a/tests/test_litellm/test_logging.py b/tests/test_litellm/test_logging.py index beba5794444..9ab362f6cd5 100644 --- a/tests/test_litellm/test_logging.py +++ b/tests/test_litellm/test_logging.py @@ -17,9 +17,15 @@ import sys import litellm from litellm._logging import ( ALL_LOGGERS, + CorrelationContextFilter, + CorrelationPlainFormatter, JsonFormatter, _initialize_loggers_with_handler, _turn_on_json, + session_id_var, + set_session_id, + set_trace_id, + trace_id_var, verbose_logger, verbose_proxy_logger, verbose_router_logger, @@ -393,3 +399,244 @@ def test_logging_calls_do_not_build_their_message_eagerly(): "these logging calls build their message eagerly; pass the values as %-style arguments instead:\n" + "\n".join(offenders) ) + + +class _JsonCapture(logging.Handler): + def __init__(self): + super().__init__() + self.formatter = JsonFormatter() + self.records: list[dict] = [] + self.addFilter(CorrelationContextFilter()) + + def emit(self, record): + self.records.append(json.loads(self.formatter.format(record))) + + +def _make_capture_logger(name: str) -> tuple[logging.Logger, _JsonCapture]: + lg = logging.getLogger(name) + cap = _JsonCapture() + lg.addHandler(cap) + lg.setLevel(logging.DEBUG) + return lg, cap + + +def test_trace_id_injected_into_json_record(monkeypatch): + """trace_id set via set_trace_id() appears in every JSON record in that context.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.trace_inject") + set_trace_id("trace-abc-123") + try: + lg.info("test message") + assert len(cap.records) == 1 + assert cap.records[0]["trace_id"] == "trace-abc-123" + finally: + trace_id_var.set("") + + +def test_session_id_injected_when_set(monkeypatch): + """session_id set via set_session_id() appears in JSON record.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.session_inject") + set_session_id("sess-xyz-456") + try: + lg.info("another message") + assert cap.records[0]["session_id"] == "sess-xyz-456" + finally: + session_id_var.set("") + + +def test_trace_id_and_session_id_cannot_be_spoofed_by_message_content(monkeypatch): + """A log message that happens to parse as JSON/dict with "trace_id"/"session_id" + keys (e.g. the proxy logging a raw request-header dict) must not override the + real correlation ids set via set_trace_id()/set_session_id().""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.spoof_attempt") + set_trace_id("real-trace-id") + set_session_id("real-session-id") + try: + lg.info('{"trace_id": "attacker-supplied-trace", "session_id": "attacker-supplied-session"}') + assert cap.records[0]["trace_id"] == "real-trace-id" + assert cap.records[0]["session_id"] == "real-session-id" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_trace_id_and_session_id_cannot_be_injected_with_no_active_context(monkeypatch): + """A message that happens to parse as JSON/dict with "trace_id"/"session_id" keys + must not surface those fields at all when CorrelationContextFilter hasn't stamped + this record - e.g. a log line emitted before Logging.__init__() runs for a request + (request_correlation_in_logs on, but no genuine trace/session id active yet).""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.no_context_spoof_attempt") + trace_id_var.set("") + session_id_var.set("") + lg.info('{"trace_id": "attacker-supplied-trace", "session_id": "attacker-supplied-session"}') + assert "trace_id" not in cap.records[0] + assert "session_id" not in cap.records[0] + + +def test_trace_id_and_session_id_are_redacted_when_credential_shaped(monkeypatch): + """A caller-controlled trace_id/session_id (e.g. from x-litellm-trace-id or a W3C + baggage header) that happens to look like a real credential must not reach log + records unredacted. CorrelationContextFilter stamps trace_id/session_id onto the + record after SecretRedactionFilter has already run, so those two fields would + otherwise bypass credential redaction entirely - the fix redacts at set_trace_id()/ + set_session_id() time instead, before the value ever reaches a log record.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_capture_logger("test.credential_shaped_correlation_id") + poisoned_trace_id = "sk-ant-api03-" + "A" * 40 + poisoned_session_id = "AKIA" + "B" * 16 + set_trace_id(poisoned_trace_id) + set_session_id(poisoned_session_id) + try: + lg.info("some benign log line") + assert cap.records[0]["trace_id"] == "REDACTED" + assert cap.records[0]["session_id"] == "REDACTED" + assert poisoned_trace_id not in json.dumps(cap.records[0]) + assert poisoned_session_id not in json.dumps(cap.records[0]) + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_session_id_absent_when_not_set(): + """session_id must NOT appear in JSON record when not set for this context.""" + lg, cap = _make_capture_logger("test.no_session") + session_id_var.set("") + lg.info("no session message") + assert "session_id" not in cap.records[0] + + +def test_trace_id_absent_when_not_set(): + """trace_id must NOT appear when not set.""" + lg, cap = _make_capture_logger("test.no_trace") + trace_id_var.set("") + lg.info("no trace message") + assert "trace_id" not in cap.records[0] + + +@pytest.mark.asyncio +async def test_contextvar_isolation_between_tasks(): + """Two concurrent async tasks each see only their own trace_id.""" + results: dict[str, str] = {} + + async def task(task_id: str, trace_id: str) -> None: + set_trace_id(trace_id) + await asyncio.sleep(0) + results[task_id] = trace_id_var.get() + + await asyncio.gather( + task("A", "trace-for-A"), + task("B", "trace-for-B"), + ) + + assert results["A"] == "trace-for-A" + assert results["B"] == "trace-for-B" + + +def test_trace_id_not_in_log_when_flag_disabled(monkeypatch): + """When request_correlation_in_logs is False (default), trace_id must not appear in JSON records even when set.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + lg, cap = _make_capture_logger("test.no_trace_gated") + set_trace_id("trace-should-not-appear") + try: + lg.info("message") + assert "trace_id" not in cap.records[0] + finally: + trace_id_var.set("") + + +def test_session_id_not_in_log_when_flag_disabled(monkeypatch): + """When request_correlation_in_logs is False (default), session_id must not appear in JSON records even when set.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + lg, cap = _make_capture_logger("test.no_session_gated") + set_session_id("sess-should-not-appear") + try: + lg.info("message") + assert "session_id" not in cap.records[0] + finally: + session_id_var.set("") + + +class _PlainCapture(logging.Handler): + def __init__(self): + super().__init__() + self.formatter = CorrelationPlainFormatter("%(message)s") + self.records: list[str] = [] + self.addFilter(CorrelationContextFilter()) + + def emit(self, record): + self.records.append(self.formatter.format(record)) + + +def _make_plain_capture_logger(name: str) -> tuple[logging.Logger, _PlainCapture]: + lg = logging.getLogger(name) + cap = _PlainCapture() + lg.addHandler(cap) + lg.setLevel(logging.DEBUG) + return lg, cap + + +def test_plain_formatter_appends_trace_id_and_session_id(monkeypatch): + """CorrelationPlainFormatter must append trace_id/session_id to non-JSON log lines too.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_plain_capture_logger("test.plain_trace_session") + set_trace_id("plain-trace-1") + set_session_id("plain-session-1") + try: + lg.info("plaintext message") + assert cap.records[0] == "plaintext message [trace_id=plain-trace-1 session_id=plain-session-1]" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_plain_formatter_appends_only_trace_id_when_session_id_absent(monkeypatch): + """Only trace_id is appended when session_id was never set.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + lg, cap = _make_plain_capture_logger("test.plain_trace_only") + set_trace_id("plain-trace-2") + session_id_var.set("") + try: + lg.info("plaintext message") + assert cap.records[0] == "plaintext message [trace_id=plain-trace-2]" + finally: + trace_id_var.set("") + + +def test_plain_formatter_unchanged_when_flag_disabled(monkeypatch): + """When request_correlation_in_logs is False, plain log lines are unmodified even if the contextvars are set.""" + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + lg, cap = _make_plain_capture_logger("test.plain_flag_off") + set_trace_id("should-not-appear") + set_session_id("should-not-appear") + try: + lg.info("plaintext message") + assert cap.records[0] == "plaintext message" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_set_trace_id_strips_control_characters(): + """set_trace_id() must strip \\r/\\n/escape sequences so a caller-controlled + trace id can't forge fake log entries when interpolated into plain-text logs.""" + token = set_trace_id('evil\r\n{"level": "CRITICAL", "message": "forged"}') + try: + value = trace_id_var.get() + assert "\r" not in value + assert "\n" not in value + finally: + trace_id_var.reset(token) + + +def test_set_session_id_bounds_length(): + """set_session_id() must bound length so an oversized caller-supplied value + isn't repeated across every log line for the request.""" + token = set_session_id("a" * 1000) + try: + assert len(session_id_var.get()) == 256 + finally: + session_id_var.reset(token) + diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 67fa827a8e4..d97d9515f08 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -6757,6 +6757,68 @@ async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error(): assert mock_create.call_args.kwargs["model"] == "owning-model" +@pytest.mark.asyncio +async def test_acreate_batch_surfaces_owning_provider_error_without_disable_fallbacks(): + """The router itself has to keep a batch inside the group that owns the input file: + the proxy only sets disable_fallbacks on the managed-files route, so the caller + otherwise gets the fallback provider's error for a file it never received.""" + from litellm.types.utils import LiteLLMBatch + + router = litellm.Router( + model_list=[ + { + "model_name": "owning-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-owning", + }, + }, + { + "model_name": "fallback-model", + "litellm_params": { + "model": "azure/gpt-4o-mini", + "api_key": "sk-fallback", + "api_base": "https://fallback.openai.azure.com", + "api_version": "2024-08-01-preview", + }, + }, + ], + fallbacks=[{"owning-model": ["fallback-model"]}], + num_retries=0, + ) + attempted_models = [] + + async def _acreate_batch(model, **kwargs): + attempted_models.append(model) + if model == "owning-model": + raise litellm.APIConnectionError( + message="Connection error - openai is unreachable", + model="openai/gpt-4o-mini", + llm_provider="openai", + ) + return LiteLLMBatch( + id="batch-created-on-the-wrong-provider", + completion_window="24h", + created_at=0, + endpoint="/v1/chat/completions", + input_file_id="file-owned-by-openai", + object="batch", + status="validating", + ) + + with patch.object(router, "_acreate_batch", _acreate_batch): + with pytest.raises(litellm.APIConnectionError, match="openai is unreachable"): + await router.acreate_batch( + model="owning-model", + input_file_id="file-owned-by-openai", + endpoint="/v1/chat/completions", + completion_window="24h", + metadata={"team": "batch-jobs"}, + ) + + assert attempted_models == ["owning-model"] + + @pytest.mark.asyncio async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): import httpx @@ -7419,6 +7481,47 @@ class TestAutoRouterMaxInputCharsWiring: assert self._registered_auto_router(router).max_input_chars == DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS +class TestGetAllowedFailsFromPolicy: + def _make_router(self, **policy_kwargs) -> litellm.Router: + from litellm.types.router import AllowedFailsPolicy + + return litellm.Router( + model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "fake"}}], + allowed_fails_policy=AllowedFailsPolicy(**policy_kwargs), + ) + + def test_no_policy_returns_none(self): + router = litellm.Router( + model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "fake"}}], + ) + assert router.get_allowed_fails_from_policy(litellm.RateLimitError("429", "openai", "gpt-4")) is None + + def test_internal_server_error_allowed_fails(self): + router = self._make_router(InternalServerErrorAllowedFails=7) + exc = litellm.InternalServerError("500", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) == 7 + + def test_service_unavailable_error_allowed_fails(self): + router = self._make_router(ServiceUnavailableErrorAllowedFails=4) + exc = litellm.ServiceUnavailableError("503", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) == 4 + + def test_bad_gateway_error_allowed_fails(self): + router = self._make_router(BadGatewayErrorAllowedFails=2) + exc = litellm.BadGatewayError("502", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) == 2 + + def test_not_found_error_allowed_fails(self): + router = self._make_router(NotFoundErrorAllowedFails=1) + exc = litellm.NotFoundError("404", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) == 1 + + def test_unmatched_exception_returns_none(self): + router = self._make_router(InternalServerErrorAllowedFails=5) + exc = litellm.RateLimitError("429", "openai", "gpt-4") + assert router.get_allowed_fails_from_policy(exc) is None + + class _LogCapture(logging.Handler): def __init__(self, level): super().__init__(level=level) @@ -7550,6 +7653,8 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): assert capture.messages, "the fallback failure path did not log at ERROR" assert huge_message not in "".join(capture.messages) assert max(len(message) for message in capture.messages) < 5_000 + + def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets(): request_kwargs = {"metadata": {}} litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7) diff --git a/tests/test_litellm/test_router_weighted_failover.py b/tests/test_litellm/test_router_weighted_failover.py index 8faf6bcd9cf..0115638e1fe 100644 --- a/tests/test_litellm/test_router_weighted_failover.py +++ b/tests/test_litellm/test_router_weighted_failover.py @@ -13,6 +13,7 @@ from unittest.mock import AsyncMock, patch import pytest +import litellm from litellm import Router from litellm.utils import _get_excluded_filtered_deployments @@ -56,16 +57,12 @@ class TestGetExcludedFilteredDeployments: # error. Returning the original list here would re-include the # just-failed deployment and let weighted failover re-pick it. deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments( - deps, excluded_deployment_ids=["a", "b"] - ) + result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["a", "b"]) assert result == [] def test_excluded_set_with_unknown_ids(self): deps = [_make_dep("a"), _make_dep("b")] - result = _get_excluded_filtered_deployments( - deps, excluded_deployment_ids=["zzz"] - ) + result = _get_excluded_filtered_deployments(deps, excluded_deployment_ids=["zzz"]) assert len(result) == 2 def test_handles_missing_model_info(self): @@ -100,6 +97,190 @@ def test_set_failed_deployment_id_on_exception(): assert exc.failed_deployment_id == "dep-a" +def test_stamp_failed_deployment_id_with_effective_model_info_prefers_kwargs(): + """kwargs["model_info"] (the dynamic client-side-credential id, when present) must win + over the static deployment's model_info, so a bad-credential tenant's failures are + attributed to their own dynamic deployment id, not the shared static one.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key"}, + "model_info": {"id": "dep-a"}, + } + ], + ) + exc = Exception("fail") + router._stamp_failed_deployment_id_with_effective_model_info( + exc, _make_dep("dep-a"), {"model_info": {"id": "dynamic-dep"}} + ) + assert exc.failed_deployment_id == "dynamic-dep" + + +def test_stamp_failed_deployment_id_with_effective_model_info_falls_back_to_deployment(): + """With no dynamic id in kwargs (the common, non-client-side-credential case), the + static deployment's own model_info.id must still be stamped.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key"}, + "model_info": {"id": "dep-a"}, + } + ], + ) + exc = Exception("fail") + router._stamp_failed_deployment_id_with_effective_model_info(exc, _make_dep("dep-a"), {}) + assert exc.failed_deployment_id == "dep-a" + + +@pytest.mark.asyncio +async def test_ageneric_api_call_with_fallbacks_helper_stamps_failed_deployment_id(): + """_ageneric_api_call_with_fallbacks_helper must stamp failed_deployment_id on a + failure, same as _completion/_acompletion, so callers identifying the failed + deployment (cooldown, weighted failover) work for this call type too instead of + depending on which metadata bucket ("metadata" vs "litellm_metadata") it uses.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, + "model_info": {"id": "dep-a"}, + } + ], + ) + + async def _failing_original_function(**kwargs): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError) as exc_info: + await router._ageneric_api_call_with_fallbacks_helper( + model="test-model", + original_generic_function=_failing_original_function, + ) + + assert getattr(exc_info.value, "failed_deployment_id", None) == "dep-a" + + +@pytest.mark.asyncio +async def test_ageneric_api_call_with_fallbacks_helper_stamps_dynamic_id_for_clientside_credentials(): + """A client-side-credential call (tenant-supplied api_key) generates a dynamic + deployment id distinct from the shared static deployment. Stamping the static id + instead would let one tenant's bad credentials cool down the deployment every + other tenant sharing this config relies on.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, + "model_info": {"id": "dep-a"}, + } + ], + ) + + async def _failing_original_function(**kwargs): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError) as exc_info: + await router._ageneric_api_call_with_fallbacks_helper( + model="test-model", + original_generic_function=_failing_original_function, + api_key="tenant-supplied-key", + litellm_metadata={"model_group": "test-model"}, + ) + + failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) + assert failed_deployment_id is not None + assert failed_deployment_id != "dep-a" + + +@pytest.mark.asyncio +async def test_acompletion_stamps_dynamic_id_for_clientside_credentials(): + """Same bug as the generic-API-call helper above, but in the regular completion + path: _acompletion's exception handlers must stamp the dynamic client-side-credential + deployment id, not the shared static deployment's id.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, + "model_info": {"id": "dep-a"}, + } + ], + ) + + with patch("litellm.acompletion", new_callable=AsyncMock, side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError) as exc_info: + await router._acompletion( + model="test-model", + messages=[{"role": "user", "content": "Hello"}], + api_key="tenant-supplied-key", + metadata={"model_group": "test-model"}, + ) + + failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) + assert failed_deployment_id is not None + assert failed_deployment_id != "dep-a" + + +@pytest.mark.asyncio +async def test_acompletion_stamps_dynamic_id_for_clientside_credentials_on_timeout(): + """Same bug as the RuntimeError case above, but for the separate `except litellm.Timeout` + branch in `_acompletion`: it has its own call to the stamping helper, so a fix that only + covers the generic `except Exception` branch would leave a caller-supplied timeout + (`litellm.Timeout` is what `x-litellm-timeout` maps to) stamping the shared static id.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, + "model_info": {"id": "dep-a"}, + } + ], + ) + + timeout_exc = litellm.Timeout(message="boom", model="test-model", llm_provider="openai") + with patch("litellm.acompletion", new_callable=AsyncMock, side_effect=timeout_exc): + with pytest.raises(litellm.Timeout) as exc_info: + await router._acompletion( + model="test-model", + messages=[{"role": "user", "content": "Hello"}], + api_key="tenant-supplied-key", + metadata={"model_group": "test-model"}, + ) + + failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) + assert failed_deployment_id is not None + assert failed_deployment_id != "dep-a" + + +def test_completion_stamps_dynamic_id_for_clientside_credentials(): + """Sync counterpart: _completion's exception handler must stamp the dynamic + client-side-credential deployment id, not the shared static deployment's id.""" + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, + "model_info": {"id": "dep-a"}, + } + ], + ) + + with patch("litellm.completion", side_effect=RuntimeError("boom")): + with pytest.raises(RuntimeError) as exc_info: + router._completion( + model="test-model", + messages=[{"role": "user", "content": "Hello"}], + api_key="tenant-supplied-key", + metadata={"model_group": "test-model"}, + ) + + failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) + assert failed_deployment_id is not None + assert failed_deployment_id != "dep-a" + + @pytest.mark.asyncio async def test_maybe_run_weighted_failover_returns_none_without_failed_id(): router = Router( @@ -641,12 +822,8 @@ async def test_maybe_run_weighted_failover_skips_when_remaining_all_in_cooldown( input_kwargs={}, ) - assert ( - result is None - ), "Should return None when all remaining deployments are in cooldown" - assert ( - not run_async_fallback_called - ), "run_async_fallback must NOT be called when no healthy deployments remain" + assert result is None, "Should return None when all remaining deployments are in cooldown" + assert not run_async_fallback_called, "run_async_fallback must NOT be called when no healthy deployments remain" @pytest.mark.asyncio @@ -705,9 +882,7 @@ async def test_maybe_run_weighted_failover_proceeds_when_one_healthy_remains( ) assert result == "ok from C" - assert ( - run_async_fallback_called - ), "run_async_fallback must be called when a healthy deployment remains" + assert run_async_fallback_called, "run_async_fallback must be called when a healthy deployment remains" @pytest.mark.asyncio diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index e80960a22c9..ed5a9f1dd63 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1,4 +1,5 @@ import json +import logging import os import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -11,6 +12,13 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm +from litellm._logging import ( + CorrelationContextFilter, + JsonFormatter, + session_id_var, + trace_id_var, + verbose_logger, +) from litellm.proxy.utils import is_valid_api_key from litellm.types.utils import ( CallTypes, @@ -21,6 +29,7 @@ from litellm.types.utils import ( StreamingChoices, Usage, ) +from litellm.types.utils import all_litellm_params from litellm.utils import ( ProviderConfigManager, TextCompletionStreamWrapper, @@ -28,6 +37,7 @@ from litellm.utils import ( _is_streaming_request, get_api_key, get_llm_provider, + get_non_default_completion_params, get_optional_params_image_gen, get_prompt_cache_min_tokens, is_cached_message, @@ -5125,3 +5135,154 @@ def test_ai21_api_key_is_resolved_from_the_documented_env_var(monkeypatch: pytes monkeypatch.setenv("AI21_API_KEY", "sk-ai21-resolved-from-env") assert get_api_key(llm_provider="ai21", dynamic_api_key=None) == "sk-ai21-resolved-from-env" + + +class _JsonCapture(logging.Handler): + def __init__(self): + super().__init__() + self.formatter = JsonFormatter() + self.records: list[dict] = [] + self.addFilter(CorrelationContextFilter()) + + def emit(self, record): + self.records.append(json.loads(self.formatter.format(record))) + + +def _make_capture_logger(name: str) -> tuple[logging.Logger, _JsonCapture]: + lg = logging.getLogger(name) + cap = _JsonCapture() + lg.addHandler(cap) + lg.setLevel(logging.DEBUG) + return lg, cap + + +@pytest.mark.asyncio +async def test_wrapper_async_restores_originating_task_context_after_success(monkeypatch): + """A successful acompletion() dispatches async_success_handler via + asyncio.create_task + the global logging worker - a different Task than the + one running acompletion() itself (this test's own task). That handler's own + restore only fixes up the detached child task it runs in; wrapper_async's own + finally block (in litellm/utils.py) must separately restore the *originating* + task's trace_id/session_id, since nothing else does. + """ + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + trace_id_var.set("outer-trace-wrapper-test") + session_id_var.set("outer-session-wrapper-test") + try: + await litellm.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + litellm_session_id="mock-call-session", + num_retries=0, + ) + assert trace_id_var.get() == "outer-trace-wrapper-test" + assert session_id_var.get() == "outer-session-wrapper-test" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_function_setup_failure_after_logging_construction_restores_context(monkeypatch): + """If function_setup() constructs Logging() (which already mutated + trace_id_var/session_id_var in __init__) but then raises before returning, + the caller's wrapper() never gets a logging_obj reference to restore from. + function_setup()'s own except block must restore the correlation context + itself in that case, or it leaks into every subsequent log line in this + thread/task until something unrelated happens to reset it.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + def _boom(self, *args, **kwargs): + raise RuntimeError("simulated failure after Logging() construction") + + monkeypatch.setattr(Logging, "update_environment_variables", _boom) + + trace_id_var.set("pre-setup-failure-trace") + session_id_var.set("pre-setup-failure-session") + try: + with pytest.raises(RuntimeError, match="simulated failure"): + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + litellm_session_id="doomed-call-session", + num_retries=0, + ) + assert trace_id_var.get() == "pre-setup-failure-trace" + assert session_id_var.get() == "pre-setup-failure-session" + finally: + trace_id_var.set("") + session_id_var.set("") + + +def test_function_setup_failure_log_line_shows_outer_not_doomed_ids(monkeypatch): + """The 'Error in function_setup' diagnostic log line itself must be stamped + with the outer/pre-call correlation ids, not the doomed call's own ids - + restoring context must happen *before* logging the exception, not after, + since the failed call never produces a usable logging object for anything + else to be attributed to.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + + def _boom(self, *args, **kwargs): + raise RuntimeError("simulated failure after Logging() construction") + + monkeypatch.setattr(Logging, "update_environment_variables", _boom) + + lg, cap = _make_capture_logger("test.function_setup_failure_log_order") + # verbose_logger is a distinct, module-level logger from our throwaway one - + # temporarily attach the same capture handler so we see its own emitted record. + verbose_logger.addHandler(cap) + try: + trace_id_var.set("outer-trace") + session_id_var.set("outer-session") + with pytest.raises(RuntimeError, match="simulated failure"): + litellm.completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + mock_response="Hello there!", + litellm_session_id="doomed-call-session", + num_retries=0, + ) + setup_failure_records = [r for r in cap.records if "Error in function_setup" in r.get("message", "")] + assert len(setup_failure_records) == 1 + record = setup_failure_records[0] + assert record.get("session_id") == "outer-session" + assert record.get("trace_id") == "outer-trace" + finally: + verbose_logger.removeHandler(cap) + trace_id_var.set("") + session_id_var.set("") + + +WEBSEARCH_INTERNAL_CONTROL_FIELDS = ( + "_websearch_interception_emit_native_blocks", + "_websearch_interception_converted_stream", +) + + +def test_websearch_interception_control_fields_never_reach_the_provider(): + """The web-search interception hooks stamp these onto kwargs to carry state + across the agentic loop. Anything the param builder does not recognize is + swept into the provider request, and a provider that validates its body + rejects the whole call: Bedrock Converse answers + `_websearch_interception_emit_native_blocks: Extra inputs are not permitted` + with a 400, so enabling interception breaks every request it touches. + + Their code-interpreter counterparts are already registered; these were not. + """ + kwargs = { + "a_real_provider_specific_param": 1, + **{field: True for field in WEBSEARCH_INTERNAL_CONTROL_FIELDS}, + } + + non_default = get_non_default_completion_params(kwargs) + + assert non_default == {"a_real_provider_specific_param": 1}, ( + "web-search interception control fields leaked into the provider params: " + f"{sorted(set(non_default) - {'a_real_provider_specific_param'})}" + ) + assert set(WEBSEARCH_INTERNAL_CONTROL_FIELDS) <= set(all_litellm_params) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index d621e85f09b..c990ae52ff2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23149 + "limit": 23057 }, "LIT002": { - "limit": 27166 + "limit": 27156 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1086 + "limit": 1078 }, "LIT007": { "limit": 0 @@ -27,9 +27,9 @@ "limit": 0 }, "LIT010": { - "limit": 16758 + "limit": 16744 }, "LIT011": { - "limit": 5598 + "limit": 5596 } } diff --git a/ui/litellm-dashboard/knip.json b/ui/litellm-dashboard/knip.json index 48b39e8122d..f6cd8ace112 100644 --- a/ui/litellm-dashboard/knip.json +++ b/ui/litellm-dashboard/knip.json @@ -1,6 +1,6 @@ { "$schema": "https://unpkg.com/knip@5/schema.json", - "entry": ["scripts/**/*.{ts,mjs}", "src/components/ui/**/*.{ts,tsx}"], + "entry": ["scripts/**/*.{ts,mjs}", "src/components/ui/**/*.{ts,tsx}", "src/**/*.test-d.{ts,tsx}"], "project": ["src/**/*.{ts,tsx}", "tests/**/*.{ts,tsx}", "scripts/**/*.{ts,mjs}"], "ignore": ["src/lib/http/schema.d.ts"], "ignoreDependencies": [ diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 3953b41a2c2..62bf5fff4b4 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -10,6 +10,7 @@ "lint": "eslint .", "test": "vitest", "test:dot": "vitest --reporter=dot", + "test:types": "vitest --run --typecheck.only", "test:watch": "vitest -w", "test:coverage": "vitest run --coverage", "format": "prettier --write .", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx index 07e5e4edf50..7d94cae468d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.test.tsx @@ -87,7 +87,7 @@ describe("CacheLeakageCard", () => { [ "Input tokens you sent in this range that weren't served from or written to the cache", "Share of your input tokens that were served from the cache", - "About how much you'd save if this uncached input used prompt caching. Estimated as uncached input tokens times the per-token discount your cached traffic already gets (realized cache savings ÷ cache-read tokens).", + "About how much you'd save if this uncached input used prompt caching. Estimated as uncached input tokens times what your cached traffic already nets per cached token (realized cache savings, after write premiums, ÷ cache read and write tokens). Blank when caching is not currently saving anything overall.", ].forEach((info) => expect(getByLabelText(info)).toBeInTheDocument()); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx index 3bc5443b2ea..3791765e11e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CacheLeakageCard.tsx @@ -107,7 +107,8 @@ const CacheLeakageCard: React.FC = ({ activity }) => { Cache leakage by {dimension === "model" ? "model" : "virtual key"}

{subject} sending large volumes of uncached input with a low cache hit rate are likely missing prompt - caching. Potential savings is approximate: uncached input priced at the realized cache-read discount. + caching. Potential savings is approximate: uncached input priced at what your cached traffic nets per + cached token, after cache-write premiums.

@@ -148,7 +149,7 @@ const CacheLeakageCard: React.FC = ({ activity }) => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx index 33c64ecf18a..c6d5a410418 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.test.tsx @@ -1,5 +1,11 @@ import { fireEvent, render } from "@testing-library/react"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const { useAuthorizedMock } = vi.hoisted(() => ({ useAuthorizedMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: useAuthorizedMock, +})); vi.mock("./UsageTab", () => ({ __esModule: true, default: () =>
})); vi.mock("./PromptCompressionTab", () => ({ __esModule: true, default: () =>
})); @@ -11,9 +17,16 @@ vi.mock("./AutoRouterBenchmarksTab", () => ({ import CostOptimizationView from "./CostOptimizationView"; -const renderView = () => render(); +const renderView = (userRole = "Admin") => { + useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole }); + return render(); +}; describe("CostOptimizationView", () => { + beforeEach(() => { + useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole: "Admin" }); + }); + it("renders the four cost-optimization tabs", () => { const { getByText } = renderView(); @@ -34,4 +47,29 @@ describe("CostOptimizationView", () => { expect(getByRole("tab", { name: "Overall" })).toHaveAttribute("aria-selected", "false"); expect(getByRole("tab", { name: "Prompt Compression" })).toHaveAttribute("aria-selected", "true"); }); + + // Unlike the other three pages in this cleanup, Cost Optimization keeps its + // nav entry for internal users: the Overall tab runs on /user/daily/activity, + // which every role may call. Only the tabs reading proxy-wide config and + // telemetry (/config/list, /auto_router/benchmarks, guardrail management) + // are proxy-admin-only, so those are what disappear. + describe("proxy-admin-only tabs", () => { + it.each(["Internal User", "Internal Viewer", "Org Admin"])("shows %s the Overall tab only", (userRole) => { + const { getByRole, queryByRole } = renderView(userRole); + + expect(getByRole("tab", { name: "Overall" })).toBeInTheDocument(); + expect(queryByRole("tab", { name: "Prompt Compression" })).not.toBeInTheDocument(); + expect(queryByRole("tab", { name: "Prompt Caching" })).not.toBeInTheDocument(); + expect(queryByRole("tab", { name: "Auto-Router" })).not.toBeInTheDocument(); + }); + + it("never mounts the panels behind the admin-only endpoints for an internal user", () => { + const { getByTestId, queryByTestId } = renderView("Internal User"); + + expect(getByTestId("usage-tab")).toBeInTheDocument(); + expect(queryByTestId("compression-tab")).not.toBeInTheDocument(); + expect(queryByTestId("caching-tab")).not.toBeInTheDocument(); + expect(queryByTestId("autorouter-benchmarks-tab")).not.toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx index 6af1e8d0441..517a0d9bd85 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/CostOptimizationView.tsx @@ -4,6 +4,7 @@ import React from "react"; import { PiggyBank } from "lucide-react"; import { Alert, Tabs } from "antd"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import UsageTab from "./UsageTab"; import PromptCompressionTab from "./PromptCompressionTab"; import PromptCachingTab from "./PromptCachingTab"; @@ -18,6 +19,7 @@ interface CostOptimizationViewProps { const CostOptimizationView: React.FC = ({ accessToken, userId, userRole }) => { const activity = useDailyActivityRange(accessToken, userId, userRole); + const canViewProxyWideCostData = useCan("viewProxyWideCostData"); const items = [ { @@ -25,21 +27,25 @@ const CostOptimizationView: React.FC = ({ accessToken label: "Overall", children: , }, - { - key: "compression", - label: "Prompt Compression", - children: , - }, - { - key: "caching", - label: "Prompt Caching", - children: , - }, - { - key: "autorouter-usage", - label: "Auto-Router", - children: , - }, + ...(canViewProxyWideCostData + ? [ + { + key: "compression", + label: "Prompt Compression", + children: , + }, + { + key: "caching", + label: "Prompt Caching", + children: , + }, + { + key: "autorouter-usage", + label: "Auto-Router", + children: , + }, + ] + : []), ]; return ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index 96d4644804b..ad68111bba7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -7,6 +7,12 @@ import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; const mockGetToolSpend = vi.fn(); +const { useAuthorizedMock } = vi.hoisted(() => ({ useAuthorizedMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: useAuthorizedMock, +})); + vi.mock("@/components/networking", () => ({ getToolSpend: (...args: unknown[]) => mockGetToolSpend(...args), })); @@ -88,11 +94,18 @@ interface RenderOptions { toolSpend?: ToolSpendResponse; from?: Date; to?: Date; + userRole?: string; } const renderWith = (results: DailyData[], options: RenderOptions = {}) => { - const { toolSpend = emptyToolSpend, from = new Date(2026, 6, 1), to = new Date(2026, 6, 14) } = options; + const { + toolSpend = emptyToolSpend, + from = new Date(2026, 6, 1), + to = new Date(2026, 6, 14), + userRole = "Admin", + } = options; mockGetToolSpend.mockResolvedValue(toolSpend); + useAuthorizedMock.mockReturnValue({ accessToken: "test-token", userId: "u1", userRole }); return render( { const toolLegends = getAllByTestId("chart-legend").filter((legend) => legend.textContent === "search,read_file"); expect(toolLegends).toHaveLength(1); }); + + // `/v1/tool/spend` is proxy-admin-only while the daily-activity charts around + // it are not, so this one card is dropped rather than the whole tab. + describe("proxy-admin-only spend-by-tool card", () => { + const toolSpend = { + by_tool: [{ tool_name: "search", spend: 4.0, call_count: 3, total_tokens: 150 }], + daily: [{ date: "2026-07-12", tool_name: "search", spend: 4.0, call_count: 3 }], + start_date: "2026-07-12", + end_date: "2026-07-12", + }; + + it.each(["Internal User", "Internal Viewer", "Org Admin"])( + "hides the card and never calls the endpoint for %s", + async (userRole) => { + const { queryByText, getByTestId } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { + toolSpend, + userRole, + }); + + // Liveness gate: the daily-activity charts still render for this role, + // so the absence below is the gate, not an empty tab. + expect(getByTestId("donut-chart")).toBeInTheDocument(); + expect(queryByText("Spend by tool")).not.toBeInTheDocument(); + await vi.waitFor(() => expect(mockGetToolSpend).not.toHaveBeenCalled()); + }, + ); + + it("keeps the card and the endpoint call for an admin", async () => { + const { findByText } = renderWith([day("2026-07-12", { compression_savings_spend: 0.04 })], { toolSpend }); + + expect(await findByText("Spend by tool")).toBeInTheDocument(); + expect(mockGetToolSpend).toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx index bd9d4f3c873..530f85dc83b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.tsx @@ -8,6 +8,7 @@ import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import { getToolSpend, ToolSpendResponse } from "@/components/networking"; import { SpendMetrics } from "@/components/UsagePage/types"; import { formatNumberWithCommas } from "@/utils/dataUtils"; @@ -82,12 +83,13 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { const startTime = dateValue.from ?? null; const endTime = dateValue.to ?? null; - const toolSpendEnabled = !!accessToken && !!startTime && !!endTime; + const canViewProxyWideCostData = useCan("viewProxyWideCostData"); + const toolSpendEnabled = canViewProxyWideCostData && !!accessToken && !!startTime && !!endTime; const rangeKey = startTime && endTime ? `${isoDay(startTime)}|${isoDay(endTime)}` : ""; const [toolSpendState, setToolSpendState] = useState<{ key: string; data: ToolSpendResponse } | null>(null); useEffect(() => { - if (!accessToken || !startTime || !endTime) return; + if (!canViewProxyWideCostData || !accessToken || !startTime || !endTime) return; let cancelled = false; getToolSpend(accessToken, isoDay(startTime), isoDay(endTime)) .then((res) => { @@ -99,7 +101,7 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { return () => { cancelled = true; }; - }, [accessToken, startTime, endTime, rangeKey]); + }, [canViewProxyWideCostData, accessToken, startTime, endTime, rangeKey]); const toolSpend = toolSpendState?.key === rangeKey ? toolSpendState.data : null; const toolSpendLoading = toolSpendEnabled && toolSpend === null; @@ -198,8 +200,8 @@ const UsageTab: React.FC = ({ accessToken, activity }) => { = ({ accessToken, activity }) => {
- - - Spend by tool -

- Spend on requests that invoked each tool (MCP and client-side tools); declaring a tool without invoking it - does not count. A request that invoked multiple tools counts its full spend toward each, so this attributes - rather than partitions spend. -

-
- - {topTools.length === 0 ? ( -

- {toolSpendLoading ? "Loading..." : "No tool usage in this range."} + {canViewProxyWideCostData && ( + + + Spend by tool +

+ Spend on requests that invoked each tool (MCP and client-side tools); declaring a tool without invoking it + does not count. A request that invoked multiple tools counts its full spend toward each, so this + attributes rather than partitions spend.

- ) : ( -
-
-

Total by tool

- + + + {topTools.length === 0 ? ( +

+ {toolSpendLoading ? "Loading..." : "No tool usage in this range."} +

+ ) : ( +
+
+

Total by tool

+ +
+
+

Daily spend by tool

+ + +
-
-

Daily spend by tool

- - -
-
- )} - - + )} + + + )}
); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts index 14fb26c53ef..0f6339f3f55 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts @@ -102,12 +102,53 @@ describe("computeCacheLeakage", () => { leaker: { alias: "leaker", metrics: { prompt_tokens: 500 } }, }), ]; - const { rows, discountPerToken } = computeCacheLeakage(results); - expect(discountPerToken).toBeCloseTo(0.002, 6); + const { rows, netSavingsPerCachedToken } = computeCacheLeakage(results); + expect(netSavingsPerCachedToken).toBeCloseTo(0.002, 6); expect(rows.map((r) => r.label)).toEqual(["leaker"]); expect(rows[0].potentialSavings).toBeCloseTo(1.0, 6); }); + it("divides net savings by cache writes as well as reads, since a new cacher pays write premiums too", () => { + const results = [ + day("2026-07-01", { + cacher: { + alias: "cacher", + metrics: { + prompt_tokens: 2000, + cache_read_input_tokens: 1000, + cache_creation_input_tokens: 1000, + prompt_caching_savings_spend: 2.0, + }, + }, + leaker: { alias: "leaker", metrics: { prompt_tokens: 500 } }, + }), + ]; + const { rows, netSavingsPerCachedToken } = computeCacheLeakage(results); + expect(netSavingsPerCachedToken).toBeCloseTo(0.001, 6); + expect(rows[0].potentialSavings).toBeCloseTo(0.5, 6); + }); + + it("declines to price leakage when write premiums leave caching net negative", () => { + const results = [ + day("2026-07-01", { + writer: { + alias: "writer", + metrics: { + prompt_tokens: 2000, + cache_read_input_tokens: 100, + cache_creation_input_tokens: 1500, + prompt_caching_savings_spend: -0.75, + }, + }, + leaker: { alias: "leaker", metrics: { prompt_tokens: 500 } }, + }), + ]; + const { rows, netSavingsPerCachedToken } = computeCacheLeakage(results); + expect(netSavingsPerCachedToken).toBeLessThan(0); + expect(rows.every((r) => r.potentialSavings === null)).toBe(true); + expect(rows.map((r) => r.label)).toEqual(["leaker", "writer"]); + }); + it("returns null estimate and ranks by uncached tokens when nobody used caching", () => { const results = [ day("2026-07-01", { @@ -115,8 +156,8 @@ describe("computeCacheLeakage", () => { small: { alias: "small", metrics: { prompt_tokens: 100 } }, }), ]; - const { rows, discountPerToken } = computeCacheLeakage(results); - expect(discountPerToken).toBeNull(); + const { rows, netSavingsPerCachedToken } = computeCacheLeakage(results); + expect(netSavingsPerCachedToken).toBeNull(); expect(rows.map((r) => r.label)).toEqual(["big", "small"]); expect(rows.every((r) => r.potentialSavings === null)).toBe(true); }); @@ -174,8 +215,8 @@ describe("computeCacheLeakage by model", () => { "claude-haiku-4-5": { prompt_tokens: 500 }, }), ]; - const { rows, discountPerToken } = computeCacheLeakage(results, "model"); - expect(discountPerToken).toBeCloseTo(0.002, 6); + const { rows, netSavingsPerCachedToken } = computeCacheLeakage(results, "model"); + expect(netSavingsPerCachedToken).toBeCloseTo(0.002, 6); expect(rows.map((r) => r.id)).toEqual(["claude-haiku-4-5"]); expect(rows[0].potentialSavings).toBeCloseTo(1.0, 6); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts index d63266c5ee7..71f9c63fe99 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts @@ -25,7 +25,7 @@ export interface CacheLeakageRow { export interface CacheLeakageResult { rows: CacheLeakageRow[]; - discountPerToken: number | null; + netSavingsPerCachedToken: number | null; } export const isAnthropicModel = (model: string): boolean => /claude|anthropic/i.test(model); @@ -97,12 +97,18 @@ export const computeCacheLeakage = ( const totals = [...byEntity.values()].reduce( (agg, a) => ({ - cacheReadTokens: agg.cacheReadTokens + a.cacheReadTokens, + cachedTokens: agg.cachedTokens + a.cacheReadTokens + a.cacheCreationTokens, realizedCachingSavings: agg.realizedCachingSavings + a.realizedCachingSavings, }), - { cacheReadTokens: 0, realizedCachingSavings: 0 }, + { cachedTokens: 0, realizedCachingSavings: 0 }, ); - const discountPerToken = totals.cacheReadTokens > 0 ? totals.realizedCachingSavings / totals.cacheReadTokens : null; + // prompt_caching_savings_spend is net of the cache-write premium, so the rate has to + // divide by every token that took the cache path -- a key that starts caching pays + // those write premiums too. Dividing by reads alone overstates it and, on write-heavy + // traffic where the net is negative, would flip the sign of a real loss into a saving + const netSavingsPerCachedToken = totals.cachedTokens > 0 ? totals.realizedCachingSavings / totals.cachedTokens : null; + // A non-positive rate prices no leakage: there is no saving to extrapolate from + const rate = netSavingsPerCachedToken != null && netSavingsPerCachedToken > 0 ? netSavingsPerCachedToken : null; const rows: CacheLeakageRow[] = [...byEntity.entries()] .map(([id, a]) => { @@ -113,18 +119,18 @@ export const computeCacheLeakage = ( sublabel: dimension === "model" ? null : a.teamId, uncachedPromptTokens, cacheHitRatio: a.promptTokens > 0 ? a.cacheReadTokens / a.promptTokens : 0, - potentialSavings: discountPerToken != null ? uncachedPromptTokens * discountPerToken : null, + potentialSavings: rate != null ? uncachedPromptTokens * rate : null, }; }) .filter((row) => row.uncachedPromptTokens > 0); const sorted = rows.sort((x, y) => - discountPerToken != null + rate != null ? (y.potentialSavings ?? 0) - (x.potentialSavings ?? 0) : y.uncachedPromptTokens - x.uncachedPromptTokens, ); - return { rows: sorted.slice(0, limit), discountPerToken }; + return { rows: sorted.slice(0, limit), netSavingsPerCachedToken }; }; export interface DailyToolSpendPoint { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/page.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/page.integration.test.tsx new file mode 100644 index 00000000000..d4c68841299 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/page.integration.test.tsx @@ -0,0 +1,53 @@ +import { screen, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import GuardrailsMonitor from "./page"; +import { renderWithProviders, testQueryClient } from "../../../../tests/test-utils"; + +const { useAuthorizedMock } = vi.hoisted(() => ({ useAuthorizedMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: useAuthorizedMock, +})); + +const fetchMock = vi.fn(); + +const requestedUrls = () => fetchMock.mock.calls.map(([url]) => String(url)); + +const renderAs = (userRole: string) => { + useAuthorizedMock.mockReturnValue({ accessToken: "sk-test", userId: "u1", userRole }); + return renderWithProviders(); +}; + +// `/guardrails/usage/*` aggregates across tenants and is listed in +// admin_viewer_routes, so it is proxy-admin-only. Nothing on this page works +// for a non-admin, hence the whole page is gated rather than a section of it. +describe("Guardrails Monitor page access by role", () => { + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + fetchMock.mockResolvedValue({ + ok: true, + status: 200, + statusText: "OK", + json: async () => ({ rows: [], chart: [], totalRequests: 0, totalBlocked: 0, passRate: 100 }), + }); + vi.stubGlobal("fetch", fetchMock); + }); + + it("fetches the guardrails usage overview for an admin", async () => { + renderAs("Admin"); + + await waitFor(() => expect(requestedUrls().some((url) => url.includes("/guardrails/usage/overview"))).toBe(true)); + }); + + it.each(["Internal User", "Internal Viewer", "Org Admin", "Unknown Role"])( + "renders the admin-only notice and fires no usage request for %s", + async (userRole) => { + renderAs(userRole); + + expect(await screen.findByText("Guardrails Monitor is only available to admin users.")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).not.toHaveBeenCalled()); + expect(requestedUrls().filter((url) => url.includes("/guardrails/usage"))).toEqual([]); + }, + ); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/page.tsx index 0c4e69c2d80..255769182bf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/page.tsx @@ -1,9 +1,17 @@ "use client"; import GuardrailsMonitorView from "./_components/GuardrailsMonitorView"; +import { AdminOnlyNotice } from "@/components/shared/AdminOnlyNotice"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import useCan from "@/app/(dashboard)/hooks/useCan"; export default function GuardrailsMonitor() { const { accessToken } = useAuthorized(); + const canViewGuardrailUsage = useCan("viewGuardrailUsage"); + + if (!canViewGuardrailUsage) { + return ; + } + return ; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts index 66dfc43cebb..fa3f15124cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts @@ -820,6 +820,18 @@ describe("useAllTeams", () => { }); const requestedPage = (url: string) => new URLSearchParams(url.split("?")[1]).get("page"); + const requestedUserId = (url: string) => new URLSearchParams(url.split("?")[1]).get("user_id"); + const asRole = (userRole: string, userId = "test-user-id") => + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId, + userRole, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); it("paginates /v2/team/list to completion and concatenates every page", async () => { fetchMock.mockImplementation((url: string) => @@ -892,4 +904,63 @@ describe("useAllTeams", () => { await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); }); + + it("scopes the request to the caller for an internal user and returns their teams", async () => { + asRole("Internal User", "member-7"); + fetchMock.mockResolvedValue(pageResponse(mockTeams, 1, 1)); + + const { result } = renderHook(() => useAllTeams(), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + expect(result.current.data).toEqual(mockTeams); + expect(result.current.data?.length).toBeGreaterThan(0); + expect(requestedUserId(fetchMock.mock.calls[0][0] as string)).toBe("member-7"); + }); + + it("carries user_id on every page of a scoped multi-page result", async () => { + asRole("Internal Viewer", "member-7"); + fetchMock.mockImplementation((url: string) => + Promise.resolve( + requestedPage(url) === "1" ? pageResponse([mockTeams[0]], 1, 2) : pageResponse([mockTeams[1]], 2, 2), + ), + ); + + const { result } = renderHook(() => useAllTeams(), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + expect(fetchMock).toHaveBeenCalledTimes(2); + const scopes = fetchMock.mock.calls.map((call) => requestedUserId(call[0] as string)); + expect(scopes).toEqual(["member-7", "member-7"]); + }); + + it.each(["Admin", "Admin Viewer", "Org Admin"])( + "sends no user_id for %s so the broad list is left intact", + async (userRole) => { + asRole(userRole); + fetchMock.mockResolvedValue(pageResponse(mockTeams, 1, 1)); + + const { result } = renderHook(() => useAllTeams(), { wrapper }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + + expect(requestedUserId(fetchMock.mock.calls[0][0] as string)).toBeNull(); + }, + ); + + it("refetches when the scope changes even though the access token has not", async () => { + asRole("Internal User", "member-7"); + fetchMock.mockResolvedValue(pageResponse(mockTeams, 1, 1)); + + const { result, rerender } = renderHook(() => useAllTeams(), { wrapper }); + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(fetchMock).toHaveBeenCalledTimes(1); + + asRole("Internal User", "member-8"); + rerender(); + + await waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + expect(requestedUserId(fetchMock.mock.calls[1][0] as string)).toBe("member-8"); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts index 4061026b94d..e209a1d7273 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts @@ -5,6 +5,7 @@ import { fetchTeams } from "@/app/(dashboard)/networking"; import { createQueryKeys } from "@/app/(dashboard)/hooks/common/queryKeysFactory"; import { teamInfoCall } from "@/components/networking"; import { getProxyBaseUrl, getGlobalLitellmHeaderName, deriveErrorMessage, handleError } from "@/components/networking"; +import { teamListScopeUserId } from "@/utils/roles"; export interface TeamsResponse { teams: Team[]; @@ -116,24 +117,30 @@ export const useTeams = (): UseQueryResult => { const ALL_TEAMS_PAGE_SIZE = 100; -const fetchAllTeamsPaged = async (accessToken: string): Promise => { - const firstPage: TeamsResponse = await teamListCall(accessToken, 1, ALL_TEAMS_PAGE_SIZE); +const fetchAllTeamsPaged = async (accessToken: string, userID: string | null): Promise => { + const firstPage: TeamsResponse = await teamListCall(accessToken, 1, ALL_TEAMS_PAGE_SIZE, { userID }); const totalPages = firstPage.total_pages ?? 1; if (totalPages <= 1) return firstPage.teams; const remainingPages: TeamsResponse[] = await Promise.all( - Array.from({ length: totalPages - 1 }, (_, i) => teamListCall(accessToken, i + 2, ALL_TEAMS_PAGE_SIZE)), + Array.from({ length: totalPages - 1 }, (_, i) => teamListCall(accessToken, i + 2, ALL_TEAMS_PAGE_SIZE, { userID })), ); return [firstPage, ...remainingPages].flatMap((page) => page.teams); }; export const useAllTeams = (): UseQueryResult => { - const { accessToken } = useAuthorized(); + const { accessToken, userId, userRole } = useAuthorized(); + const scopedUserID = teamListScopeUserId(userRole, userId); return useQuery({ queryKey: teamKeys.list({ - filters: { scope: "all", pageSize: ALL_TEAMS_PAGE_SIZE, accessToken: accessToken ?? "" }, + filters: { + scope: "all", + pageSize: ALL_TEAMS_PAGE_SIZE, + accessToken: accessToken ?? "", + userID: scopedUserID ?? "", + }, }), - queryFn: async () => await fetchAllTeamsPaged(accessToken!), + queryFn: async () => await fetchAllTeamsPaged(accessToken!, scopedUserID), enabled: Boolean(accessToken), staleTime: 30000, }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.test.ts new file mode 100644 index 00000000000..2215817b618 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.test.ts @@ -0,0 +1,135 @@ +import { getUiSettings } from "@/components/networking"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { PTU_FLAG_REFRESH_MS, usePtuCostAttributionEnabled } from "./usePtuCostAttributionEnabled"; +import { useUISettings } from "./useUISettings"; + +vi.mock("@/components/networking", () => ({ + getUiSettings: vi.fn(), +})); + +describe("usePtuCostAttributionEnabled", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + vi.clearAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + /** Read the flag alongside the query it derives from, so assertions wait for a settled fetch. */ + const renderSettledFlag = async (settings: unknown) => { + (getUiSettings as any).mockResolvedValue(settings); + const { result } = renderHook(() => ({ enabled: usePtuCostAttributionEnabled(), query: useUISettings() }), { + wrapper, + }); + await waitFor(() => { + expect(result.current.query.isSuccess).toBe(true); + }); + return result; + }; + + it("is true only when the proxy reports the flag as enabled", async () => { + const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: true } }); + expect(result.current.enabled).toBe(true); + }); + + it("is false when the proxy reports the flag as disabled", async () => { + const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: false } }); + expect(result.current.enabled).toBe(false); + }); + + it("is false when the proxy omits the flag entirely", async () => { + const result = await renderSettledFlag({ values: { enable_chat_ui: true } }); + expect(result.current.enabled).toBe(false); + }); + + it("is false when the proxy returns no values at all", async () => { + const result = await renderSettledFlag({}); + expect(result.current.enabled).toBe(false); + }); + + it("does not treat a truthy non-boolean as enabled", async () => { + const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: "false" } }); + expect(result.current.enabled).toBe(false); + }); + + it("does not treat the string 'true' as enabled, since the proxy sends a real boolean", async () => { + const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: "true" } }); + expect(result.current.enabled).toBe(false); + }); + + it("is false before the settings request resolves", () => { + (getUiSettings as any).mockReturnValue(new Promise(() => {})); + const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper }); + expect(result.current).toBe(false); + }); + + it("is false when the settings request fails", async () => { + (getUiSettings as any).mockRejectedValue(new Error("boom")); + const { result } = renderHook(() => ({ enabled: usePtuCostAttributionEnabled(), query: useUISettings() }), { + wrapper, + }); + await waitFor(() => { + expect(result.current.query.isError).toBe(true); + }); + expect(result.current.enabled).toBe(false); + }); +}); + +describe("staleness", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + vi.clearAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("polls the flag once it is on, so an already-open dashboard notices it going off", async () => { + (getUiSettings as any).mockResolvedValue({ values: { enable_ptu_cost_attribution: true } }); + const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper }); + await waitFor(() => { + expect(result.current).toBe(true); + }); + + const observers = queryClient.getQueryCache().getAll()[0].observers; + const polling = observers.filter((o: any) => o.options.refetchInterval === PTU_FLAG_REFRESH_MS); + expect(polling.length).toBeGreaterThan(0); + expect(polling[0].options.staleTime).toBe(PTU_FLAG_REFRESH_MS); + expect(PTU_FLAG_REFRESH_MS).toBeLessThan(60 * 60 * 1000); + }); + + it("does not poll while the flag is off, which is every deployment that never opted in", async () => { + // The hook cannot gate on the flag before reading it, so it starts on the shared + // one-hour cache and only escalates once it has seen the feature enabled. Polling + // unconditionally made a disabled deployment re-fetch settings 120x more often. + (getUiSettings as any).mockResolvedValue({ values: { enable_ptu_cost_attribution: false } }); + const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper }); + await waitFor(() => { + expect(result.current).toBe(false); + }); + + const observers = queryClient.getQueryCache().getAll()[0].observers; + expect(observers.every((o: any) => o.options.refetchInterval === undefined)).toBe(true); + expect(observers.every((o: any) => o.options.staleTime === 60 * 60 * 1000)).toBe(true); + }); + + it("leaves the default alone for every other settings consumer", async () => { + (getUiSettings as any).mockResolvedValue({ values: {} }); + const { result } = renderHook(() => useUISettings(), { wrapper }); + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + const observers = queryClient.getQueryCache().getAll()[0].observers; + expect(observers[0].options.staleTime).toBe(60 * 60 * 1000); + expect(observers[0].options.refetchInterval).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.ts new file mode 100644 index 00000000000..e9b5afac562 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.ts @@ -0,0 +1,26 @@ +import { useUISettings } from "./useUISettings"; + +export const PTU_COST_ATTRIBUTION_SETTING_KEY = "enable_ptu_cost_attribution"; + +/** + * Whether the proxy opted into PTU flat-cost attribution. + * + * Derived on the proxy from LITELLM_ENABLE_PTU_COST_ATTRIBUTION and returned read-only on + * /get/ui_settings, so it is not editable from the UI. Anything other than an explicit + * true (including a settings fetch that has not resolved) counts as off. + * + * Polled only once the flag has been seen on. This tracks the proxy process rather than a + * persisted setting, so an already-open dashboard has to notice a restart that turns the + * feature off, and a form that stays mounted and focused never refetches on staleTime + * alone. A deployment that never opts in is the common case and gets the shared one-hour + * cache, so the poll costs nothing where the feature is unused; the trade is that turning + * it on reaches an open dashboard on the next natural refetch rather than within 30s. + */ +export const PTU_FLAG_REFRESH_MS = 30 * 1000; + +export const usePtuCostAttributionEnabled = (): boolean => { + const { data } = useUISettings(); + const enabled = data?.values?.[PTU_COST_ATTRIBUTION_SETTING_KEY] === true; + useUISettings(enabled ? { staleTime: PTU_FLAG_REFRESH_MS, refetchInterval: PTU_FLAG_REFRESH_MS } : undefined); + return enabled; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts index 14c6c5e3888..749fc98c0d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts @@ -4,11 +4,21 @@ import { createQueryKeys } from "../common/queryKeysFactory"; const uiSettingsKeys = createQueryKeys("uiSettings"); -export const useUISettings = () => { +/** + * UI settings, cached for an hour by default because they rarely change. + * + * Both options are per observer in react-query, so a caller reading a value that tracks + * proxy process state, rather than a persisted setting, can refresh it on its own cadence + * without changing how long every other caller caches. `staleTime` alone only marks the + * cached copy stale; a screen that stays mounted and focused never refetches on its own, + * so a caller that needs to notice a change also has to poll. + */ +export const useUISettings = (options?: { staleTime?: number; refetchInterval?: number }) => { return useQuery>({ queryKey: uiSettingsKeys.list({}), queryFn: async () => await getUiSettings(), - staleTime: 60 * 60 * 1000, // 1 hour - data rarely changes + staleTime: options?.staleTime ?? 60 * 60 * 1000, // 1 hour - data rarely changes gcTime: 60 * 60 * 1000, // 1 hour - keep in cache for 1 hour + refetchInterval: options?.refetchInterval, }); }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useCan.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useCan.ts index f538e1dff15..13903007cae 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useCan.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useCan.ts @@ -3,10 +3,12 @@ import { hasCapability, type Capability } from "@/utils/capabilities"; import useAuthorized from "./useAuthorized"; +import useIsOrgAdmin from "./useIsOrgAdmin"; const useCan = (capability: Capability): boolean => { const { userRole } = useAuthorized(); - return hasCapability(userRole, capability); + const isOrgAdmin = useIsOrgAdmin(); + return hasCapability(userRole, capability, isOrgAdmin); }; export default useCan; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useIsOrgAdmin.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useIsOrgAdmin.test.ts new file mode 100644 index 00000000000..bb913e2fd64 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useIsOrgAdmin.test.ts @@ -0,0 +1,53 @@ +/* @vitest-environment jsdom */ +import { renderHook } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { Organization } from "@/components/networking"; +import useIsOrgAdmin from "./useIsOrgAdmin"; + +const { useAuthorizedMock, useOrganizationsMock } = vi.hoisted(() => ({ + useAuthorizedMock: vi.fn(), + useOrganizationsMock: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: useAuthorizedMock })); +vi.mock("./organizations/useOrganizations", () => ({ useOrganizations: useOrganizationsMock })); + +const orgWithMembers = (members: { user_id: string; user_role: string }[]): Organization => + ({ organization_id: "org-1", members }) as unknown as Organization; + +const renderAs = (userRole: string, organizations: Organization[] | undefined) => { + useAuthorizedMock.mockReturnValue({ userId: "user-1", userRole }); + useOrganizationsMock.mockReturnValue({ data: organizations }); + return renderHook(() => useIsOrgAdmin()).result; +}; + +describe("useIsOrgAdmin", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("is true for the session a real org admin carries: internal_user plus an org_admin membership", () => { + const result = renderAs("Internal User", [orgWithMembers([{ user_id: "user-1", user_role: "org_admin" }])]); + expect(result.current).toBe(true); + }); + + it("is false for an internal user with no org_admin membership", () => { + const result = renderAs("Internal User", [orgWithMembers([{ user_id: "user-1", user_role: "internal_user" }])]); + expect(result.current).toBe(false); + }); + + it("is false while the organization list is still loading", () => { + const result = renderAs("Internal User", undefined); + expect(result.current).toBe(false); + }); + + it("is true for a session role of org_admin even with no membership rows", () => { + expect(renderAs("org_admin", []).current).toBe(true); + expect(renderAs("Org Admin", []).current).toBe(true); + }); + + it("is false for a proxy admin, who is covered by role-based gates instead", () => { + const result = renderAs("Admin", []); + expect(result.current).toBe(false); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useIsOrgAdmin.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useIsOrgAdmin.ts new file mode 100644 index 00000000000..d93b57a3fc3 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useIsOrgAdmin.ts @@ -0,0 +1,14 @@ +"use client"; + +import { isOrgAdminForAnyOrg, isOrgAdminSessionRole } from "@/utils/roles"; + +import { useOrganizations } from "./organizations/useOrganizations"; +import useAuthorized from "./useAuthorized"; + +const useIsOrgAdmin = (): boolean => { + const { userId, userRole } = useAuthorized(); + const { data: organizations } = useOrganizations(); + return isOrgAdminSessionRole(userRole) || isOrgAdminForAnyOrg(organizations, userId); +}; + +export default useIsOrgAdmin; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/page.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/page.integration.test.tsx new file mode 100644 index 00000000000..8d15bb59187 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/page.integration.test.tsx @@ -0,0 +1,59 @@ +import { screen, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import Memory from "./page"; +import { renderWithProviders, testQueryClient } from "../../../../tests/test-utils"; + +const { useAuthorizedMock } = vi.hoisted(() => ({ useAuthorizedMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: useAuthorizedMock, +})); + +const fetchMock = vi.fn(); + +const requestedUrls = () => fetchMock.mock.calls.map(([url]) => String(url)); + +const renderAs = (userRole: string) => { + useAuthorizedMock.mockReturnValue({ accessToken: "sk-test", userId: "u1", userRole }); + return renderWithProviders(); +}; + +// `/v1/memory` scopes rows per caller in the handler, but the route gate keeps +// it proxy-admin-only, so a non-admin deep-linking to /ui/memory gets a 401. +describe("Memory page access by role", () => { + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + fetchMock.mockResolvedValue({ + ok: true, + status: 200, + statusText: "OK", + text: async () => "", + json: async () => ({ memories: [], total: 0 }), + }); + vi.stubGlobal("fetch", fetchMock); + }); + + it("lists memory entries for an admin", async () => { + renderAs("Admin"); + + await waitFor(() => expect(requestedUrls().some((url) => url.includes("/v1/memory"))).toBe(true)); + }); + + it.each(["Internal User", "Internal Viewer", "Org Admin", "Unknown Role"])( + "renders the admin-only notice and fires no memory request for %s", + async (userRole) => { + renderAs(userRole); + + expect(await screen.findByText("Memory is only available to admin users.")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).not.toHaveBeenCalled()); + expect(requestedUrls().filter((url) => url.includes("/v1/memory"))).toEqual([]); + }, + ); + + it("hides the deprecation banner along with the page body for a denied role", () => { + renderAs("Internal User"); + + expect(screen.queryByText(/draft deprecation list/i)).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/page.tsx index b88996c5396..7b1b6223372 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/memory/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/page.tsx @@ -2,10 +2,18 @@ import { MemoryView } from "./_components/MemoryView"; import { DeprecationBanner } from "@/components/DeprecationBanner"; +import { AdminOnlyNotice } from "@/components/shared/AdminOnlyNotice"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import useCan from "@/app/(dashboard)/hooks/useCan"; export default function Memory() { const { accessToken, userRole, userId } = useAuthorized(); + const canViewMemory = useCan("viewMemory"); + + if (!canViewMemory) { + return ; + } + return ( <> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx index e3db50b7300..0e4455ee912 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx @@ -1,6 +1,6 @@ import React from "react"; import { describe, it, expect, vi, beforeEach } from "vitest"; -import { screen, waitFor, within } from "@testing-library/react"; +import { act, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { renderWithProviders } from "../../../../../tests/test-utils"; import UsagePage from "./usage"; @@ -49,6 +49,14 @@ const renderUsage = (overrides: Partial> />, ); +// Width of this window is guarded by "proves the flush window is wide enough". +const flushPendingRequests = async () => { + await act(async () => { + await new Promise((resolve) => setTimeout(resolve, 0)); + await new Promise((resolve) => setTimeout(resolve, 0)); + }); +}; + beforeEach(() => { vi.clearAllMocks(); networking.getProxyUISettings.mockResolvedValue(UNLIMITED_SETTINGS); @@ -185,18 +193,65 @@ describe("old usage page", () => { }); }); - describe("as a non-admin", () => { - it("renders only the All Up tab and skips admin-only queries", async () => { - renderUsage({ userRole: "Internal User" }); + // org_admin is an organization membership role; those users reach the UI as "Internal User". + describe.each(["Internal User", "Internal Viewer", "internal_user", "internal_user_viewer", "Org Admin"])( + "as %s", + (userRole) => { + it("shows the admin-only notice instead of the usage dashboard", async () => { + renderUsage({ userRole }); + + expect(await screen.findByText(/Proxy-wide usage is only available to admin users/i)).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "All Up" })).not.toBeInTheDocument(); + }); + + it("fires no /global/spend or /global/activity request", async () => { + renderUsage({ userRole }); + + await screen.findByText(/Proxy-wide usage is only available to admin users/i); + await flushPendingRequests(); + + expect(networking.getProxyUISettings).not.toHaveBeenCalled(); + expect(networking.adminSpendLogsCall).not.toHaveBeenCalled(); + expect(networking.adminTopKeysCall).not.toHaveBeenCalled(); + expect(networking.adminTopModelsCall).not.toHaveBeenCalled(); + expect(networking.adminTopEndUsersCall).not.toHaveBeenCalled(); + expect(networking.teamSpendLogsCall).not.toHaveBeenCalled(); + expect(networking.tagsSpendLogsCall).not.toHaveBeenCalled(); + expect(networking.allTagNamesCall).not.toHaveBeenCalled(); + expect(networking.adminspendByProvider).not.toHaveBeenCalled(); + expect(networking.adminGlobalActivity).not.toHaveBeenCalled(); + expect(networking.adminGlobalActivityPerModel).not.toHaveBeenCalled(); + }); + }, + ); + + describe("the admin-only gate", () => { + it("proves the flush window is wide enough to catch a leaked request", async () => { + renderUsage({ userRole: "Admin" }); + + await flushPendingRequests(); + + expect(networking.getProxyUISettings).toHaveBeenCalled(); + expect(networking.adminSpendLogsCall).toHaveBeenCalled(); + expect(networking.tagsSpendLogsCall).toHaveBeenCalled(); + expect(networking.adminGlobalActivity).toHaveBeenCalled(); + }); + + it("still lets an admin through, so the notice is a real gate and not a dead branch", async () => { + renderUsage({ userRole: "Admin" }); expect(await screen.findByRole("tab", { name: "All Up" })).toBeInTheDocument(); - expect(screen.queryByRole("tab", { name: "Team Based Usage" })).not.toBeInTheDocument(); - expect(screen.queryByRole("tab", { name: "Customer Usage" })).not.toBeInTheDocument(); - expect(screen.queryByRole("tab", { name: "Tag Based Usage" })).not.toBeInTheDocument(); - + expect(screen.queryByText(/Proxy-wide usage is only available to admin users/i)).not.toBeInTheDocument(); await waitFor(() => expect(networking.adminSpendLogsCall).toHaveBeenCalled()); - expect(networking.teamSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.adminTopEndUsersCall).not.toHaveBeenCalled(); + }); + + it("does not put the session token in the provider spend query", async () => { + renderUsage({ userRole: "Admin", token: "session-jwt-value" }); + + await waitFor(() => expect(networking.adminspendByProvider).toHaveBeenCalled()); + const callArgs = networking.adminspendByProvider.mock.calls[0]; + expect(callArgs).not.toContain("session-jwt-value"); + expect(callArgs[0]).toBe("sk-test"); }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx index 3d55f9bb698..5b2f8547822 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx @@ -37,6 +37,7 @@ import { } from "@/components/networking"; import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView"; import { MoneyCell } from "@/components/shared/table_cells"; +import { hasCapability } from "@/utils/capabilities"; import { formatNumberWithCommas } from "@/utils/dataUtils"; interface UsagePageProps { @@ -90,6 +91,7 @@ const TeamSpendBarList: React.FC<{ data: TeamSpendTotal[] }> = ({ data }) => { }; const UsagePage: React.FC = ({ accessToken, token, userRole, userID, keys, premiumUser }) => { + const canViewGlobalSpend = hasCapability(userRole, "viewGlobalSpend"); const currentDate = new Date(); const [keySpendData, setKeySpendData] = useState([]); const [topKeys, setTopKeys] = useState([]); @@ -155,8 +157,11 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use }; useEffect(() => { + if (!canViewGlobalSpend) { + return; + } updateTagSpendData(dateValue.from, dateValue.to); - }, [dateValue, selectedTags]); + }, [canViewGlobalSpend, dateValue, selectedTags]); const updateEndUserData = async ( startTime: Date | undefined, @@ -319,10 +324,7 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use const fetchProviderSpend = () => fetchAndSetData( - () => - accessToken && token - ? adminspendByProvider(accessToken, token, startTime, endTime) - : Promise.reject("No access token or token"), + () => (accessToken ? adminspendByProvider(accessToken, startTime, endTime) : Promise.reject("No access token")), setSpendByProvider, "Error fetching provider spend", ); @@ -467,6 +469,9 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use useEffect(() => { const initlizeUsageData = async () => { + if (!canViewGlobalSpend) { + return; + } if (accessToken && token && userRole && userID) { const proxy_settings: ProxySettings | undefined = await fetchProxySettings(); if (proxy_settings) { @@ -493,7 +498,24 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use }; initlizeUsageData(); - }, [accessToken, token, userRole, userID, startTime, endTime]); + }, [canViewGlobalSpend, accessToken, token, userRole, userID, startTime, endTime]); + + if (!canViewGlobalSpend) { + return ( +
+ + + Usage + + +

+ Proxy-wide usage is only available to admin users. Your own usage is on the Usage page. +

+
+
+
+ ); + } if (proxySettings?.DISABLE_EXPENSIVE_DB_QUERIES) { return ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx index 3c94977f0dc..0c487af213e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx @@ -1,4 +1,5 @@ -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; +import { renderWithProviders as render } from "@/../tests/test-utils"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ChatUI from "./ChatUI"; import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index 57ff7906eda..0241ef8a77e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -25,6 +25,7 @@ import React, { useEffect, useRef, useState } from "react"; import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; import { v4 as uuidv4 } from "uuid"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; import PolicySelector from "@/components/policies/PolicySelector"; import MCPToolArgumentsForm, { MCPToolArgumentsFormRef } from "@/components/mcp_tools/MCPToolArgumentsForm"; @@ -106,6 +107,7 @@ const ChatUI: React.FC = ({ simplified = false, fixedModel, }) => { + const canViewPolicies = useCan("viewPolicies"); const [mcpServers, setMCPServers] = useState([]); const [mcpToolsets, setMCPToolsets] = useState([]); const [isToolsetsInfoModalVisible, setIsToolsetsInfoModalVisible] = useState(false); @@ -1652,32 +1654,34 @@ const ChatUI: React.FC = ({ />
-
- - Policies - - Select policy/policies to apply to this LLM API call. Policies define which guardrails are - applied based on conditions. You can set up your policies{" "} - - here - - . - - } - > - - - - -
+ {canViewPolicies && ( +
+ + Policies + + Select policy/policies to apply to this LLM API call. Policies define which guardrails are + applied based on conditions. You can set up your policies{" "} + + here + + . + + } + > + + + + +
+ )} {/* Code Interpreter Toggle - Only for Responses endpoint */} {endpointType === EndpointType.RESPONSES && ( diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx index 39346105f2a..c3b417987e6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx @@ -6,6 +6,7 @@ import { type ComplianceFramework, type CompliancePrompt, } from "@/data/compliancePrompts"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import { getGuardrailsList, testPoliciesAndGuardrails } from "@/components/networking"; import PolicySelector, { getPolicyOptionEntries } from "@/components/policies/PolicySelector"; import { Policy } from "@/components/policies/types"; @@ -123,6 +124,7 @@ export default function ComplianceUI({ fixedModel, proxySettings, }: ComplianceUIProps) { + const canViewPolicies = useCan("viewPolicies"); const frameworks = getFrameworks(); const [policyValueToLabel, setPolicyValueToLabel] = useState>(new Map()); @@ -701,29 +703,37 @@ export default function ComplianceUI({

Test Configuration

-

Select policies, guardrails, or both to test against.

+

+ {canViewPolicies + ? "Select policies, guardrails, or both to test against." + : "Select guardrails to test against."} +

-
- - {accessToken && ( - - )} -
+ {canViewPolicies && ( + <> +
+ + {accessToken && ( + + )} +
-
-
- or -
-
+
+
+ or +
+
+ + )}
{capitalizedEntityLabel} Spend Overview - - Total Spend - - ${formatNumberWithCommas(spendData.metadata.total_spend, 2)} - - - - Total Requests - {spendData.metadata.total_api_requests.toLocaleString()} - - - Successful Requests - - {spendData.metadata.total_successful_requests.toLocaleString()} - - - - Failed Requests - - {spendData.metadata.total_failed_requests.toLocaleString()} - - - - Total Tokens - {spendData.metadata.total_tokens.toLocaleString()} - + {summaryTiles.map(renderSummaryTile)} @@ -456,21 +313,40 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti new Date(a.date).getTime() - new Date(b.date).getTime())} + data={[...spendData.results] + .sort((a, b) => new Date(a.date).getTime() - new Date(b.date).getTime()) + .map((row) => ({ + ...row, + "Request cost": row.metrics.spend ?? 0, + "Flat cost": row.metrics.flat_cost ?? 0, + }))} index="date" - categories={["metrics.spend"]} - colors={["cyan"]} + categories={showFlatCost ? ["Request cost", "Flat cost"] : ["metrics.spend"]} + colors={showFlatCost ? ["cyan", "violet"] : ["cyan"]} + stack={showFlatCost} valueFormatter={valueFormatterSpend} yAxisWidth={100} - showLegend={false} + showLegend={showFlatCost} customTooltip={({ payload, active }) => { if (!active || !payload?.[0]) return null; const data = payload[0].payload; const entityCount = Object.keys(data.breakdown.entities || {}).length; + const requestSpend = data.metrics.spend ?? 0; + const flatCost = data.metrics.flat_cost ?? 0; return (

{data.date}

-

Total Spend: ${formatNumberWithCommas(data.metrics.spend, 2)}

+ {showFlatCost ? ( + <> +

Request cost: ${formatNumberWithCommas(requestSpend, 2)}

+

Flat cost: ${formatNumberWithCommas(flatCost, 2)}

+

+ Total cost: ${formatNumberWithCommas(requestSpend + flatCost, 2)} +

+ + ) : ( +

Total Spend: ${formatNumberWithCommas(data.metrics.spend, 2)}

+ )}

Total Requests: {data.metrics.api_requests}

Successful: {data.metrics.successful_requests}

Failed: {data.metrics.failed_requests}

@@ -597,7 +473,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti Top Virtual Keys = ({ accessToken, entityType, enti
- {/* Top Agents - only for team entity type */} - {entityType === "team" && ( + {showAgentBreakdown && (
Top Agents Driving Spend @@ -644,7 +519,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti `$${formatNumberWithCommas(value, 2)}`} @@ -666,7 +541,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti - {getProviderSpend().map((provider) => ( + {getProviderSpend(spendData.results).map((provider) => (
@@ -708,7 +583,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti ), }, - ...(entityType === "team" + ...(showAgentBreakdown ? [{ key: "agents", label: "Agent Activity", content: }] : []), { @@ -757,7 +632,7 @@ const EntityUsage: React.FC = ({ accessToken, entityType, enti } /> )} - {agentIsFetchingMore && entityType === "team" && ( + {agentIsFetchingMore && showAgentBreakdown && ( = ({ accessToken, entityType, enti } /> )} - {agentCancelled && entityType === "team" && ( + {agentCancelled && showAgentBreakdown && ( { + const modelSpend: { [key: string]: any } = {}; + results.forEach((day) => { + Object.entries(day.breakdown[modelBreakdownKey] || {}).forEach(([model, metrics]) => { + if (!modelSpend[model]) { + modelSpend[model] = { + spend: 0, + requests: 0, + successful_requests: 0, + failed_requests: 0, + tokens: 0, + }; + } + try { + modelSpend[model].spend += metrics.metrics.spend; + } catch (e) { + console.error(`Error adding spend for ${model}: ${e}, got metrics: ${JSON.stringify(metrics)}`); + } + modelSpend[model].requests += metrics.metrics.api_requests; + modelSpend[model].successful_requests += metrics.metrics.successful_requests; + modelSpend[model].failed_requests += metrics.metrics.failed_requests; + modelSpend[model].tokens += metrics.metrics.total_tokens; + }); + }); + + return Object.entries(modelSpend) + .map(([model, metrics]) => ({ + key: model, + ...metrics, + })) + .sort((a, b) => b.spend - a.spend) + .slice(0, topModelsLimit); +}; + +export const getTopAgents = (results: ExtendedDailyData[], topAgentsLimit: number) => { + const agentSpend: { [key: string]: any } = {}; + results.forEach((day) => { + Object.entries(day.breakdown.entities || {}).forEach(([agentId, data]) => { + if (!agentSpend[agentId]) { + agentSpend[agentId] = { + spend: 0, + requests: 0, + successful_requests: 0, + failed_requests: 0, + tokens: 0, + agent_name: (data.metadata as any)?.agent_name || agentId, + }; + } + agentSpend[agentId].spend += data.metrics.spend; + agentSpend[agentId].requests += data.metrics.api_requests; + agentSpend[agentId].successful_requests += data.metrics.successful_requests; + agentSpend[agentId].failed_requests += data.metrics.failed_requests; + agentSpend[agentId].tokens += data.metrics.total_tokens; + }); + }); + + return Object.entries(agentSpend) + .map(([agentId, metrics]) => ({ + key: metrics.agent_name, + ...metrics, + })) + .sort((a, b) => b.spend - a.spend) + .slice(0, topAgentsLimit); +}; + +export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number) => { + const keySpend: { [key: string]: KeyMetricWithMetadata } = {}; + results.forEach((day) => { + const { breakdown } = day; + const { entities } = breakdown; + const tagDictionary = Object.keys(entities).reduce((acc: { [key: string]: TagUsage[] }, entity) => { + const { api_key_breakdown } = entities[entity]; + Object.keys(api_key_breakdown).forEach((key) => { + const tagUsage = { tag: entity, usage: api_key_breakdown[key].metrics.spend }; + if (acc[key]) { + acc[key].push(tagUsage); + } else { + acc[key] = [tagUsage]; + } + }); + return acc; + }, {}); + Object.entries(day.breakdown.api_keys || {}).forEach(([key, metrics]) => { + if (!keySpend[key]) { + keySpend[key] = { + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: { + key_alias: metrics.metadata.key_alias, + team_id: metrics.metadata.team_id || null, + tags: tagDictionary[key] || [], + }, + }; + } + keySpend[key].metrics.spend += metrics.metrics.spend; + keySpend[key].metrics.prompt_tokens += metrics.metrics.prompt_tokens; + keySpend[key].metrics.completion_tokens += metrics.metrics.completion_tokens; + keySpend[key].metrics.total_tokens += metrics.metrics.total_tokens; + keySpend[key].metrics.api_requests += metrics.metrics.api_requests; + keySpend[key].metrics.successful_requests += metrics.metrics.successful_requests; + keySpend[key].metrics.failed_requests += metrics.metrics.failed_requests; + keySpend[key].metrics.cache_read_input_tokens += metrics.metrics.cache_read_input_tokens || 0; + keySpend[key].metrics.cache_creation_input_tokens += metrics.metrics.cache_creation_input_tokens || 0; + }); + }); + + return Object.entries(keySpend) + .map(([api_key, metrics]) => ({ + api_key, + key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias + tags: metrics.metadata.tags || "-", + spend: metrics.metrics.spend, + })) + .sort((a, b) => b.spend - a.spend) + .slice(0, topKeysLimit); +}; + +export const getProviderSpend = (results: ExtendedDailyData[]) => { + const providerSpend: { [key: string]: any } = {}; + results.forEach((day) => { + Object.entries(day.breakdown.providers || {}).forEach(([provider, metrics]) => { + if (!providerSpend[provider]) { + providerSpend[provider] = { + provider, + spend: 0, + requests: 0, + successful_requests: 0, + failed_requests: 0, + tokens: 0, + }; + } + try { + providerSpend[provider].spend += metrics.metrics.spend; + providerSpend[provider].requests += metrics.metrics.api_requests; + providerSpend[provider].successful_requests += metrics.metrics.successful_requests; + providerSpend[provider].failed_requests += metrics.metrics.failed_requests; + providerSpend[provider].tokens += metrics.metrics.total_tokens; + } catch (e) { + console.error(`Error processing provider ${provider}: ${e}`); + } + }); + }); + + return Object.values(providerSpend) + .filter((provider) => provider.spend > 0) + .sort((a, b) => b.spend - a.spend); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.test.ts new file mode 100644 index 00000000000..403f9473391 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.test.ts @@ -0,0 +1,82 @@ +import { describe, expect, it } from "vitest"; +import { buildCostBreakdownTiles, buildSummaryTiles, hasFlatCost } from "./entityUsageSummary"; + +const metadata = { + total_spend: 100, + total_flat_cost: 40, + total_api_requests: 12, + total_successful_requests: 10, + total_failed_requests: 2, + total_tokens: 3456, +}; + +describe("hasFlatCost", () => { + it("is false when there is no flat cost to report", () => { + expect(hasFlatCost({ ...metadata, total_flat_cost: 0 })).toBe(false); + const { total_flat_cost, ...noFlat } = metadata; + expect(hasFlatCost(noFlat)).toBe(false); + }); + + it("is true once a flat cost has accrued", () => { + expect(hasFlatCost(metadata)).toBe(true); + }); +}); + +describe("buildSummaryTiles", () => { + it("keeps the row at five tiles either way so adding flat cost never narrows the cards", () => { + expect(buildSummaryTiles(metadata, false)).toHaveLength(5); + expect(buildSummaryTiles(metadata, true)).toHaveLength(5); + }); + + it("shows request-only spend under the original title when there is no flat cost", () => { + const [first] = buildSummaryTiles(metadata, false); + expect(first.title).toBe("Total Spend"); + expect(first.value).toBe("$100.00"); + expect(first.expandable).toBeUndefined(); + }); + + it("rolls flat cost into a single expandable Total Cost tile", () => { + const [first] = buildSummaryTiles(metadata, true); + expect(first.title).toBe("Total Cost"); + expect(first.value).toBe("$140.00"); + expect(first.expandable).toBe(true); + expect(first.tooltip).toBeTruthy(); + }); + + it("never renders the breakdown titles in the top row", () => { + const titles = buildSummaryTiles(metadata, true).map((t) => t.title); + expect(titles).not.toContain("Flat Cost"); + expect(titles).not.toContain("Request Cost"); + }); + + it("treats a missing flat cost as zero", () => { + const { total_flat_cost, ...noFlat } = metadata; + expect(buildSummaryTiles(noFlat, true)[0].value).toBe("$100.00"); + }); +}); + +describe("buildCostBreakdownTiles", () => { + it("splits the total into request cost and flat cost", () => { + const byTitle = Object.fromEntries(buildCostBreakdownTiles(metadata).map((t) => [t.title, t.value])); + expect(byTitle["Request Cost"]).toBe("$100.00"); + expect(byTitle["Flat Cost"]).toBe("$40.00"); + }); + + it("adds up to the Total Cost tile so the expanded view reconciles", () => { + const parse = (v: string) => Number(v.replace(/[$,]/g, "")); + const parts = buildCostBreakdownTiles(metadata).map((t) => parse(t.value)); + expect(parts[0] + parts[1]).toBe(parse(buildSummaryTiles(metadata, true)[0].value)); + }); + + it("explains each part, including that flat cost is outside budgets", () => { + const byTitle = Object.fromEntries(buildCostBreakdownTiles(metadata).map((t) => [t.title, t.tooltip])); + expect(byTitle["Request Cost"]).toBeTruthy(); + expect(byTitle["Flat Cost"]).toContain("budget"); + }); + + it("treats a missing flat cost as zero", () => { + const { total_flat_cost, ...noFlat } = metadata; + const byTitle = Object.fromEntries(buildCostBreakdownTiles(noFlat).map((t) => [t.title, t.value])); + expect(byTitle["Flat Cost"]).toBe("$0.00"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.ts new file mode 100644 index 00000000000..093cd9c40af --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.ts @@ -0,0 +1,66 @@ +import { formatNumberWithCommas } from "@/utils/dataUtils"; + +export interface SummaryTile { + title: string; + value: string; + className?: string; + tooltip?: string; + expandable?: boolean; +} + +interface SpendSummaryMetadata { + total_spend: number; + total_flat_cost?: number; + total_api_requests: number; + total_successful_requests: number; + total_failed_requests: number; + total_tokens: number; +} + +export const TOTAL_COST_TOOLTIP = + "Request cost plus flat cost for reserved capacity. Select this tile to see the breakdown."; + +export const REQUEST_COST_TOOLTIP = + "Usage-based cost of the requests this entity sent during the selected period, priced per token."; + +export const FLAT_COST_TOOLTIP = + "Reserved provisioned throughput, billed per hour whether or not requests are sent. Reported here only; it does not count toward team, key, user, or organization budgets."; + +export const hasFlatCost = (metadata: SpendSummaryMetadata): boolean => (metadata.total_flat_cost ?? 0) > 0; + +export const buildSummaryTiles = (metadata: SpendSummaryMetadata, showFlatCost: boolean): SummaryTile[] => { + const flatCost = metadata.total_flat_cost ?? 0; + return [ + showFlatCost + ? { + title: "Total Cost", + value: `$${formatNumberWithCommas(metadata.total_spend + flatCost, 2)}`, + tooltip: TOTAL_COST_TOOLTIP, + expandable: true, + } + : { title: "Total Spend", value: `$${formatNumberWithCommas(metadata.total_spend, 2)}` }, + { title: "Total Requests", value: metadata.total_api_requests.toLocaleString() }, + { + title: "Successful Requests", + value: metadata.total_successful_requests.toLocaleString(), + className: "text-green-600", + }, + { title: "Failed Requests", value: metadata.total_failed_requests.toLocaleString(), className: "text-red-600" }, + { title: "Total Tokens", value: metadata.total_tokens.toLocaleString() }, + ]; +}; + +export const buildCostBreakdownTiles = (metadata: SpendSummaryMetadata): SummaryTile[] => [ + { + title: "Request Cost", + value: `$${formatNumberWithCommas(metadata.total_spend, 2)}`, + className: "text-cyan-600", + tooltip: REQUEST_COST_TOOLTIP, + }, + { + title: "Flat Cost", + value: `$${formatNumberWithCommas(metadata.total_flat_cost ?? 0, 2)}`, + className: "text-violet-600", + tooltip: FLAT_COST_TOOLTIP, + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx index 0ded7d195d0..9085cf961a9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx @@ -502,6 +502,8 @@ describe("UsagePage", () => { userId: "user-123", userEmail: "test@example.com", userRole: "Internal User", + userRoleLabel: "Internal User", + isViewOnly: false, premiumUser: true, disabledPersonalKeyCreation: false, showSSOBanner: false, @@ -861,6 +863,27 @@ describe("UsagePage", () => { }); }); + it.each(["organization", "agent"])("should not render the %s usage view for an internal user", async (usageView) => { + mockUseAuthorized.mockReturnValue(nonAdminSession); + + renderWithProviders(); + + await waitFor(() => { + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); + }); + + const usageSelect = screen.getByTestId("usage-view-select"); + act(() => { + fireEvent.change(usageSelect, { target: { value: "team" } }); + }); + expect(screen.getAllByText("Entity Usage").length).toBeGreaterThan(0); + + act(() => { + fireEvent.change(usageSelect, { target: { value: usageView } }); + }); + expect(screen.queryByText("Entity Usage")).not.toBeInTheDocument(); + }); + describe("admin user selector", () => { it("should render user selector for admin users in global view", async () => { renderWithProviders(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index c3645d6371e..494df313ac0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -33,6 +33,7 @@ import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers"; +import { hasCapability } from "@/utils/capabilities"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { all_admin_roles, internalUserRoles } from "@/utils/roles"; import { ActivityMetrics, processActivityData } from "@/components/activity_metrics"; @@ -109,6 +110,8 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const { data: currentUser } = useCurrentUser(); const isAdmin = all_admin_roles.includes(userRole || ""); const canViewTagUsage = isAdmin || internalUserRoles.includes(userRole || ""); + const canViewOrganizationUsage = hasCapability(userRole, "viewOrganizationUsage"); + const canViewAgentUsage = hasCapability(userRole, "viewAgentUsage"); // Debounced search for user selector const [userSearchInput, setUserSearchInput] = useState(""); @@ -513,7 +516,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { setUsageView(value)} - isAdmin={isAdmin} + userRole={userRole} canViewTagUsage={canViewTagUsage} /> @@ -950,7 +953,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { )} {/* Organization Usage Panel */} - {usageView === "organization" && ( + {usageView === "organization" && canViewOrganizationUsage && ( = ({ teams, organizations }) => { /> )} - {usageView === "agent" && ( + {usageView === "agent" && canViewAgentUsage && ( { }); it("should render", () => { - render(); + render(); expect(screen.getByText("Usage View")).toBeInTheDocument(); expect(screen.getByText("Select the usage data you want to view")).toBeInTheDocument(); expect(screen.getByRole("combobox")).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "Your Usage" })).toBeInTheDocument(); }); it("should call onChange when value changes", () => { - render(); + render(); const select = screen.getByRole("combobox"); act(() => { @@ -109,14 +110,32 @@ describe("UsageViewSelect", () => { }); it("should show Tag Usage for non-admin users with tag usage permission", () => { - render(); + render(); expect(screen.getByRole("option", { name: "Tag Usage" })).toBeInTheDocument(); }); it("should hide Tag Usage for non-admin users without tag usage permission", () => { - render(); + render(); expect(screen.queryByRole("option", { name: "Tag Usage" })).not.toBeInTheDocument(); }); + + it.each(["Organization Usage", "Agent Usage (A2A)"])("should show %s to an admin", (optionName) => { + render(); + + expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument(); + }); + + it.each(["Organization Usage", "Agent Usage (A2A)"])("should hide %s from an internal user", (optionName) => { + render(); + + expect(screen.queryByRole("option", { name: optionName })).not.toBeInTheDocument(); + }); + + it.each(["Team Usage", "Tag Usage"])("should keep %s available to an internal user", (optionName) => { + render(); + + expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx index 94b483cb539..54c1d5ab7cc 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx @@ -11,6 +11,8 @@ import { } from "@ant-design/icons"; import { Badge, Select } from "antd"; import React from "react"; +import { hasCapability, type Capability } from "@/utils/capabilities"; +import { all_admin_roles } from "@/utils/roles"; export type UsageOption = | "global" | "my-usage" @@ -24,7 +26,7 @@ export type UsageOption = export interface UsageViewSelectProps { value: UsageOption; onChange: (value: UsageOption) => void; - isAdmin: boolean; + userRole: string | null; canViewTagUsage?: boolean; title?: string; description?: string; @@ -35,6 +37,7 @@ interface OptionConfig { label: string; description: string; icon: React.ReactNode; + capability?: Capability; adminOnly?: boolean; showForAdmin?: string; showForNonAdmin?: string; @@ -63,12 +66,9 @@ const OPTIONS: OptionConfig[] = [ { value: "organization", label: "Organization Usage", - showForAdmin: "Organization Usage", - showForNonAdmin: "Your Organization Usage", - description: "View organization-level usage", - descriptionForAdmin: "View usage across all organizations", - descriptionForNonAdmin: "View your organization's usage", + description: "View usage across all organizations", icon: , + capability: "viewOrganizationUsage", }, { value: "team", @@ -95,7 +95,7 @@ const OPTIONS: OptionConfig[] = [ label: "Agent Usage (A2A)", description: "View usage by AI agents", icon: , - adminOnly: true, + capability: "viewAgentUsage", }, { value: "user", @@ -115,14 +115,18 @@ const OPTIONS: OptionConfig[] = [ export const UsageViewSelect: React.FC = ({ value, onChange, - isAdmin, + userRole, canViewTagUsage = false, title = "Usage View", description = "Select the usage data you want to view", "data-id": dataId, }) => { + const isAdmin = all_admin_roles.includes(userRole ?? ""); const getFilteredOptions = () => { return OPTIONS.filter((option) => { + if (option.capability) { + return hasCapability(userRole, option.capability); + } if (option.value === "tag" && canViewTagUsage) { return true; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.test.ts new file mode 100644 index 00000000000..b1d467074f6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.test.ts @@ -0,0 +1,51 @@ +import { describe, expect, it } from "vitest"; +import { sumMetadata } from "./usePaginatedDailyActivity"; + +describe("sumMetadata", () => { + it("sums flat cost across pages instead of keeping the first page's value", () => { + // A team whose activity spans more than one page accrues flat cost on each of them. + // Keeping page 1's value under-reports the Flat Cost and Total Cost tiles. + const merged = sumMetadata({ total_spend: 1, total_flat_cost: 174.5 }, { total_spend: 2, total_flat_cost: 777 }); + + expect(merged.total_flat_cost).toBe(951.5); + expect(merged.total_spend).toBe(3); + }); + + it("treats a page missing the field as zero rather than dropping the running total", () => { + expect(sumMetadata({ total_flat_cost: 480 }, {}).total_flat_cost).toBe(480); + expect(sumMetadata({}, { total_flat_cost: 480 }).total_flat_cost).toBe(480); + }); + + it("carries non-summable keys through from the first page", () => { + const merged = sumMetadata( + { page: 1, total_pages: 3, total_spend: 1 }, + { page: 2, total_pages: 3, total_spend: 2 }, + ); + + expect(merged.page).toBe(1); + expect(merged.total_pages).toBe(3); + }); + + it("sums every total_* metric the daily activity metadata exposes", () => { + // Guards the class of bug rather than one field: a new backend total that nobody adds + // to SUMMABLE_METADATA_KEYS freezes at page 1, and spend still looks right so it reads + // as trustworthy. + const page = { + total_spend: 1, + total_prompt_tokens: 1, + total_completion_tokens: 1, + total_tokens: 1, + total_api_requests: 1, + total_successful_requests: 1, + total_failed_requests: 1, + total_cache_read_input_tokens: 1, + total_cache_creation_input_tokens: 1, + total_flat_cost: 1, + }; + const merged = sumMetadata(page, page); + + for (const key of Object.keys(page)) { + expect(merged[key], `${key} must be summed across pages`).toBe(2); + } + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts index 86b26bda539..a8bfee5be2f 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts @@ -23,6 +23,7 @@ const SUMMABLE_METADATA_KEYS = [ "total_failed_requests", "total_cache_read_input_tokens", "total_cache_creation_input_tokens", + "total_flat_cost", ] as const; interface DailyActivityResponse { @@ -68,7 +69,12 @@ const EMPTY_DATA: DailyActivityResponse = { }, }; -function sumMetadata(a: Record, b: Record): Record { +/** + * Combine two pages of metadata. Only keys in SUMMABLE_METADATA_KEYS are added; anything + * else keeps the first page's value, so a total the backend adds later is silently frozen + * at page 1 until it is listed above. Exported so that contract can be tested directly. + */ +export function sumMetadata(a: Record, b: Record): Record { const result = { ...a }; for (const key of SUMMABLE_METADATA_KEYS) { result[key] = (a[key] || 0) + (b[key] || 0); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx index 5fcc55c1e98..42f21cd7b69 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx @@ -1,10 +1,11 @@ /* @vitest-environment jsdom */ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render, screen, waitFor } from "@testing-library/react"; +import { screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../../tests/test-utils"; import ViewUserDashboard from "./view_users"; const userListCall = vi.fn(); @@ -78,7 +79,7 @@ const defaultProps = { }; const renderDashboard = () => - render( + renderWithProviders( , diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 2c1d28d82f8..9eb7645fb2e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -1,4 +1,5 @@ import { Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; +import { parseAsString, useQueryState } from "nuqs"; import React, { useCallback, useEffect, useMemo, useState } from "react"; import { Button } from "antd"; @@ -72,7 +73,7 @@ const ViewUserDashboard: React.FC = ({ const [selectionMode, setSelectionMode] = useState(false); const [isBulkEditModalVisible, setIsBulkEditModalVisible] = useState(false); - const [selectedUserId, setSelectedUserId] = useState(null); + const [selectedUserId, setSelectedUserId] = useQueryState("user", parseAsString.withOptions({ history: "push" })); const [openInEditMode, setOpenInEditMode] = useState(false); const [editModalVisible, setEditModalVisible] = useState(false); @@ -139,15 +140,18 @@ const ViewUserDashboard: React.FC = ({ setRowSelection({}); }, []); - const handleUserClick = useCallback((userId: string, openInEdit: boolean = false) => { - setSelectedUserId(userId); - setOpenInEditMode(openInEdit); - }, []); + const handleUserClick = useCallback( + (userId: string, openInEdit: boolean = false) => { + void setSelectedUserId(userId); + setOpenInEditMode(openInEdit); + }, + [setSelectedUserId], + ); const handleCloseUserInfo = useCallback(() => { - setSelectedUserId(null); + void setSelectedUserId(null); setOpenInEditMode(false); - }, []); + }, [setSelectedUserId]); const handleDelete = useCallback((user: UserInfo) => { setUserToDelete(user); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTab.tsx new file mode 100644 index 00000000000..bf53433adea --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTab.tsx @@ -0,0 +1,102 @@ +"use client"; + +import React, { useCallback, useEffect, useMemo, useState } from "react"; + +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { indexesListCall } from "@/components/networking"; +import { VectorStore } from "@/components/vector_store_management/types"; + +import IndexesTable from "./IndexesTable"; + +export interface VectorStoreIndex { + id: string; + index_name: string; + litellm_params: { + vector_store_index: string; + vector_store_name: string; + }; + index_info?: Record | null; + created_at?: string | null; + created_by?: string | null; + updated_at?: string | null; + updated_by?: string | null; +} + +interface IndexesTabProps { + accessToken: string | null; + vectorStores: VectorStore[]; + onViewVectorStore: (vectorStoreId: string) => void; +} + +const IndexesTab: React.FC = ({ accessToken, vectorStores, onViewVectorStore }) => { + const [indexes, setIndexes] = useState([]); + const [isLoading, setIsLoading] = useState(true); + + const vectorStoreIdsByName = useMemo( + () => + new Map( + vectorStores.flatMap((store) => + store.vector_store_name ? [[store.vector_store_name, store.vector_store_id] as const] : [], + ), + ), + [vectorStores], + ); + + const resolveVectorStoreId = useCallback((name: string) => vectorStoreIdsByName.get(name), [vectorStoreIdsByName]); + + useEffect(() => { + const fetchIndexes = async () => { + if (!accessToken) { + setIsLoading(false); + return; + } + try { + const response = await indexesListCall(accessToken); + setIndexes(response.data || []); + } catch (error) { + console.error("Error fetching indexes:", error); + NotificationsManager.fromBackend("Error fetching indexes: " + error); + } finally { + setIsLoading(false); + } + }; + fetchIndexes(); + }, [accessToken]); + + return ( +
+

+ Vector store indexes registered on this proxy via the /v1/indexes API. See the{" "} + + vector store index docs + {" "} + for how this works. Index passthrough is supported for Azure AI Search and Milvus today; support for more + providers can be added, so please{" "} + + file a GitHub issue + {" "} + if you want your provider supported. +

+
+ +
+
+ ); +}; + +export default IndexesTab; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.test.tsx new file mode 100644 index 00000000000..f31c51f3fba --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.test.tsx @@ -0,0 +1,100 @@ +import { render, screen, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import type { VectorStoreIndex } from "./IndexesTab"; +import IndexesTable from "./IndexesTable"; + +vi.mock("next/navigation", async () => ({ + ...(await vi.importActual("next/navigation")), + useRouter: () => ({ push: vi.fn() }), +})); + +const newerIndex: VectorStoreIndex = { + id: "idx-newer", + index_name: "newer-index", + litellm_params: { vector_store_index: "provider-newer", vector_store_name: "newer-store" }, + created_by: "admin@example.com", + created_at: "2024-02-20T10:30:00Z", +}; + +const olderIndex: VectorStoreIndex = { + id: "idx-older", + index_name: "older-index", + litellm_params: { vector_store_index: "provider-older", vector_store_name: "older-store" }, + created_by: "admin@example.com", + created_at: "2024-01-10T09:15:00Z", +}; + +const undatedIndex: VectorStoreIndex = { + id: "idx-undated", + index_name: "undated-index", + litellm_params: { vector_store_index: "provider-undated", vector_store_name: "undated-store" }, + created_by: null, + created_at: null, +}; + +const noResolve = () => undefined; + +describe("IndexesTable", () => { + it("should display the empty state when no indexes are registered", () => { + render(); + expect(screen.getByText("No indexes registered yet")).toBeInTheDocument(); + }); + + it("should render index rows with dash fallbacks for missing created_by and created_at", () => { + render( + , + ); + expect(screen.getByText("newer-index")).toBeInTheDocument(); + expect(screen.getByText("newer-store")).toBeInTheDocument(); + expect(screen.getByText("provider-newer")).toBeInTheDocument(); + expect(screen.getByText("admin@example.com")).toBeInTheDocument(); + const undatedRow = screen.getByText("undated-index").closest("tr"); + expect(undatedRow).not.toBeNull(); + expect(within(undatedRow as HTMLElement).getAllByText("-")).toHaveLength(2); + }); + + it("should sort by created_at descending by default", () => { + render( + , + ); + const rows = screen.getAllByRole("row").slice(1); + expect(within(rows[0]).getByText("newer-index")).toBeInTheDocument(); + expect(within(rows[1]).getByText("older-index")).toBeInTheDocument(); + }); + + it("should call onViewVectorStore with the resolved id when the vector store cell is clicked", async () => { + const user = userEvent.setup(); + const onViewVectorStore = vi.fn(); + render( + (name === "newer-store" ? "vs-newer" : undefined)} + onViewVectorStore={onViewVectorStore} + />, + ); + await user.click(screen.getByRole("button", { name: "newer-store" })); + expect(onViewVectorStore).toHaveBeenCalledWith("vs-newer"); + }); + + it("should render an unresolvable vector store name as plain text without a clickable cell", () => { + render(); + expect(screen.getByText("newer-store")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "newer-store" })).not.toBeInTheDocument(); + }); + + it("should link created_by to the user detail deep link", () => { + render(); + const link = screen.getByRole("link", { name: "admin@example.com" }); + expect(link.getAttribute("href")).toMatch(/\/users\?user=admin%40example\.com$/); + }); + + it("should keep the dash fallback and render no link for a null created_by", () => { + render(); + const row = screen.getByText("undated-index").closest("tr"); + expect(row).not.toBeNull(); + expect(within(row as HTMLElement).queryByRole("link")).not.toBeInTheDocument(); + expect(within(row as HTMLElement).getAllByText("-").length).toBeGreaterThan(0); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.tsx new file mode 100644 index 00000000000..927fd48acb6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.tsx @@ -0,0 +1,62 @@ +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; + +import type { VectorStoreIndex } from "./IndexesTab"; +import { getIndexesTableColumns } from "./IndexesTableColumns"; + +interface IndexesTableProps { + data: VectorStoreIndex[]; + resolveVectorStoreId: (name: string) => string | undefined; + onViewVectorStore: (vectorStoreId: string) => void; + isLoading?: boolean; +} + +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +function EmptyState() { + return ( +
+
+ +
+
No indexes registered yet
+
Indexes registered on this proxy will appear here.
+
+ ); +} + +const IndexesTable: React.FC = ({ + data, + resolveVectorStoreId, + onViewVectorStore, + isLoading = false, +}) => { + const [sorting, setSorting] = useState(DEFAULT_SORTING); + + const columns = useMemo( + () => getIndexesTableColumns({ resolveVectorStoreId, onViewVectorStore }), + [resolveVectorStoreId, onViewVectorStore], + ); + + return ( + row.id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading indexes…" + noDataMessage={} + size="compact" + /> + ); +}; + +export default IndexesTable; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTableColumns.tsx new file mode 100644 index 00000000000..c21665f87cd --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTableColumns.tsx @@ -0,0 +1,108 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell } from "@/components/shared/table_cells"; +import { userDetailHref } from "@/utils/entityLinks"; + +import type { VectorStoreIndex } from "./IndexesTab"; + +interface IndexesTableColumnsDeps { + resolveVectorStoreId: (name: string) => string | undefined; + onViewVectorStore: (vectorStoreId: string) => void; +} + +export const getIndexesTableColumns = ({ + resolveVectorStoreId, + onViewVectorStore, +}: IndexesTableColumnsDeps): ColumnDef[] => [ + { + id: "index_name", + accessorKey: "index_name", + meta: { title: "Index Name" }, + header: ({ column }) => , + size: 220, + enableSorting: true, + cell: ({ row }) => ( + + {row.original.index_name || "-"} + + ), + }, + { + id: "vector_store_name", + accessorFn: (row) => row.litellm_params.vector_store_name, + meta: { title: "Vector Store" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => { + const name = row.original.litellm_params.vector_store_name; + const vectorStoreId = name ? resolveVectorStoreId(name) : undefined; + if (vectorStoreId) { + return ( + onViewVectorStore(vectorStoreId)} + /> + ); + } + return ( + + {name || "-"} + + ); + }, + }, + { + id: "vector_store_index", + accessorFn: (row) => row.litellm_params.vector_store_index, + meta: { title: "Provider Index" }, + header: "Provider Index", + size: 220, + enableSorting: false, + cell: ({ row }) => ( + + {row.original.litellm_params.vector_store_index || "-"} + + ), + }, + { + id: "created_by", + accessorKey: "created_by", + meta: { title: "Created By" }, + header: "Created By", + size: 160, + enableSorting: false, + cell: ({ row }) => { + const createdBy = row.original.created_by; + if (createdBy) { + return ( + + ); + } + return -; + }, + }, + { + id: "created_at", + accessorKey: "created_at", + sortingFn: "datetime", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx index 521c1f879ee..7137da201d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.test.tsx @@ -1,8 +1,8 @@ -import { render, screen } from "@testing-library/react"; +import { render, screen, waitFor } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; -import { vectorStoreListCall } from "@/components/networking"; +import { credentialListCall, indexesListCall, vectorStoreListCall } from "@/components/networking"; import VectorStoreManagement from "./index"; @@ -10,6 +10,7 @@ vi.mock("@/components/networking", () => ({ vectorStoreListCall: vi.fn(), vectorStoreDeleteCall: vi.fn(), credentialListCall: vi.fn(), + indexesListCall: vi.fn(), })); vi.mock("./VectorStoreTable", () => ({ @@ -20,11 +21,18 @@ vi.mock("./VectorStoreTable", () => ({ })); vi.mock("./VectorStoreForm", () => ({ __esModule: true, default: () => null })); -vi.mock("./vector_store_info", () => ({ __esModule: true, default: () => null })); +vi.mock("./vector_store_info", () => ({ + __esModule: true, + default: ({ vectorStoreId }: { vectorStoreId: string }) => ( +
{vectorStoreId}
+ ), +})); vi.mock("./CreateVectorStore", () => ({ __esModule: true, default: () => null })); vi.mock("./TestVectorStoreTab", () => ({ __esModule: true, default: () => null })); const mockVectorStoreListCall = vi.mocked(vectorStoreListCall); +const mockCredentialListCall = vi.mocked(credentialListCall); +const mockIndexesListCall = vi.mocked(indexesListCall); const openManageTab = async (user: ReturnType) => { await user.click(screen.getByRole("tab", { name: "Manage Vector Stores" })); @@ -60,3 +68,90 @@ describe("VectorStoreManagement loading state", () => { expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test"); }); }); + +describe("VectorStoreManagement Indexes tab", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockVectorStoreListCall.mockResolvedValue({ data: [] }); + mockCredentialListCall.mockResolvedValue({ credentials: [] }); + }); + + it("should render fetched indexes for a proxy admin after the Indexes tab is clicked", async () => { + const user = userEvent.setup(); + mockIndexesListCall.mockResolvedValue({ + object: "list", + data: [ + { + id: "idx-1", + index_name: "support-docs-index", + litellm_params: { vector_store_index: "pinecone-support-docs", vector_store_name: "support-docs-store" }, + }, + ], + }); + render(); + await user.click(screen.getByRole("tab", { name: "Indexes" })); + expect(await screen.findByText("support-docs-index")).toBeInTheDocument(); + expect(screen.getByText("support-docs-store")).toBeInTheDocument(); + expect(mockIndexesListCall).toHaveBeenCalledWith("sk-test"); + }); + + it("should not render the Indexes tab for an Admin Viewer", async () => { + render(); + await waitFor(() => expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test")); + expect(screen.getByRole("tab", { name: "Manage Vector Stores" })).toBeInTheDocument(); + expect(screen.queryByRole("tab", { name: "Indexes" })).not.toBeInTheDocument(); + }); + + it("should swap to the vector store info view when an index's vector store name is clicked", async () => { + const user = userEvent.setup(); + mockVectorStoreListCall.mockResolvedValue({ + data: [ + { + vector_store_id: "vs-1", + vector_store_name: "support-docs-store", + custom_llm_provider: "bedrock", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + ], + }); + mockIndexesListCall.mockResolvedValue({ + object: "list", + data: [ + { + id: "idx-1", + index_name: "support-docs-index", + litellm_params: { vector_store_index: "pinecone-support-docs", vector_store_name: "support-docs-store" }, + }, + ], + }); + render(); + await user.click(screen.getByRole("tab", { name: "Indexes" })); + await user.click(await screen.findByRole("button", { name: "support-docs-store" })); + expect(await screen.findByTestId("vector-store-info-view")).toHaveTextContent("vs-1"); + expect(screen.queryByText("Vector Store Management")).not.toBeInTheDocument(); + }); + + it("should link to the feature docs and a GitHub issue for unsupported providers on the Indexes tab", async () => { + const user = userEvent.setup(); + mockIndexesListCall.mockResolvedValue({ object: "list", data: [] }); + render(); + await user.click(screen.getByRole("tab", { name: "Indexes" })); + expect(screen.getByRole("link", { name: "vector store index docs" })).toHaveAttribute( + "href", + "https://docs.litellm.ai/docs/providers/azure_ai/azure_ai_vector_stores_passthrough", + ); + expect(screen.getByRole("link", { name: "file a GitHub issue" })).toHaveAttribute( + "href", + "https://github.com/BerriAI/litellm/issues", + ); + expect(screen.getByText(/supported for Azure AI Search and Milvus today/)).toBeInTheDocument(); + }); + + it("should not call indexesListCall until the Indexes tab is clicked", async () => { + render(); + await waitFor(() => expect(mockVectorStoreListCall).toHaveBeenCalledWith("sk-test")); + expect(screen.getByRole("tab", { name: "Indexes" })).toBeInTheDocument(); + expect(mockIndexesListCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx index 8afa49ea148..1e526131fa3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx @@ -13,7 +13,8 @@ import DeleteResourceModal from "@/components/common_components/DeleteResourceMo import VectorStoreInfoView from "./vector_store_info"; import CreateVectorStore from "./CreateVectorStore"; import TestVectorStoreTab from "./TestVectorStoreTab"; -import { isAdminRole } from "@/utils/roles"; +import IndexesTab from "./IndexesTab"; +import { isAdminRole, isProxyAdminRole } from "@/utils/roles"; import NotificationsManager from "@/components/molecules/notifications_manager"; import { Button } from "@/components/ui/button"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -163,6 +164,11 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID Test Vector Store + {isProxyAdminRole(userRole || "") && ( + + Indexes + + )} @@ -188,6 +194,12 @@ const VectorStoreManagement: React.FC = ({ accessToken, userID + + {isProxyAdminRole(userRole || "") && ( + + + + )} {/* Create Vector Store Modal */} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/vector_store_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/vector_store_info.test.tsx new file mode 100644 index 00000000000..b5dc359b328 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/vector_store_info.test.tsx @@ -0,0 +1,81 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { credentialListCall, vectorStoreInfoCall } from "@/components/networking"; + +import VectorStoreInfoView from "./vector_store_info"; + +vi.mock("@/components/networking", () => ({ + vectorStoreInfoCall: vi.fn(), + vectorStoreUpdateCall: vi.fn(), + credentialListCall: vi.fn(), +})); + +vi.mock("./VectorStoreTester", () => ({ __esModule: true, default: () => null })); + +const mockVectorStoreInfoCall = vi.mocked(vectorStoreInfoCall); +const mockCredentialListCall = vi.mocked(credentialListCall); + +describe("VectorStoreInfoView", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockCredentialListCall.mockResolvedValue({ credentials: [] }); + }); + + it("should render the store details once the fetch resolves", async () => { + mockVectorStoreInfoCall.mockResolvedValue({ + vector_store: { + vector_store_id: "vs-1", + vector_store_name: "support-docs-store", + custom_llm_provider: "bedrock", + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + }, + }); + render( + , + ); + expect(await screen.findByText("Vector Store ID: vs-1")).toBeInTheDocument(); + }); + + it("should show a not-found state with a working back button when the fetch fails instead of loading forever", async () => { + const user = userEvent.setup(); + const onClose = vi.fn(); + mockVectorStoreInfoCall.mockRejectedValue(new Error("Vector store not found")); + render( + , + ); + expect(await screen.findByText("Vector store not found")).toBeInTheDocument(); + expect(screen.getByText(/vs-gone could not be loaded/)).toBeInTheDocument(); + expect(screen.queryByText("Loading...")).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: /Back to Vector Stores/ })); + expect(onClose).toHaveBeenCalled(); + }); + + it("should show the not-found state when the fetch resolves without a vector store", async () => { + mockVectorStoreInfoCall.mockResolvedValue({ vector_store: null }); + render( + , + ); + expect(await screen.findByText("Vector store not found")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/vector_store_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/vector_store_info.tsx index e5646037d14..4f94a7d56f6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/vector_store_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/vector_store_info.tsx @@ -33,6 +33,7 @@ const VectorStoreInfoView: React.FC = ({ }) => { const [form] = Form.useForm(); const [vectorStoreDetails, setVectorStoreDetails] = useState(null); + const [loadFailed, setLoadFailed] = useState(false); const [isEditing, setIsEditing] = useState(editVectorStore); const [metadataString, setMetadataString] = useState("{}"); const [credentials, setCredentials] = useState([]); @@ -40,31 +41,35 @@ const VectorStoreInfoView: React.FC = ({ const fetchVectorStoreDetails = async () => { if (!accessToken) return; try { + setLoadFailed(false); const response = await vectorStoreInfoCall(accessToken, vectorStoreId); - if (response && response.vector_store) { - setVectorStoreDetails(response.vector_store); + if (!response || !response.vector_store) { + setLoadFailed(true); + return; + } + setVectorStoreDetails(response.vector_store); - // If metadata exists and is an object, stringify it for display/editing - if (response.vector_store.vector_store_metadata) { - const metadata = - typeof response.vector_store.vector_store_metadata === "string" - ? JSON.parse(response.vector_store.vector_store_metadata) - : response.vector_store.vector_store_metadata; - setMetadataString(JSON.stringify(metadata, null, 2)); - } + // If metadata exists and is an object, stringify it for display/editing + if (response.vector_store.vector_store_metadata) { + const metadata = + typeof response.vector_store.vector_store_metadata === "string" + ? JSON.parse(response.vector_store.vector_store_metadata) + : response.vector_store.vector_store_metadata; + setMetadataString(JSON.stringify(metadata, null, 2)); + } - if (editVectorStore) { - form.setFieldsValue({ - vector_store_id: response.vector_store.vector_store_id, - custom_llm_provider: response.vector_store.custom_llm_provider, - vector_store_name: response.vector_store.vector_store_name, - vector_store_description: response.vector_store.vector_store_description, - }); - } + if (editVectorStore) { + form.setFieldsValue({ + vector_store_id: response.vector_store.vector_store_id, + custom_llm_provider: response.vector_store.custom_llm_provider, + vector_store_name: response.vector_store.vector_store_name, + vector_store_description: response.vector_store.vector_store_description, + }); } } catch (error) { console.error("Error fetching vector store details:", error); NotificationsManager.fromBackend("Error fetching vector store details: " + error); + setLoadFailed(true); } }; @@ -113,6 +118,20 @@ const VectorStoreInfoView: React.FC = ({ } }; + if (loadFailed) { + return ( +
+ + Vector store not found + + Vector store {vectorStoreId} could not be loaded. It may have been deleted. + +
+ ); + } + if (!vectorStoreDetails) { return
Loading...
; } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/workflows/page.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/workflows/page.integration.test.tsx new file mode 100644 index 00000000000..6b332faf704 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/workflows/page.integration.test.tsx @@ -0,0 +1,59 @@ +import { screen, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import Workflows from "./page"; +import { renderWithProviders, testQueryClient } from "../../../../tests/test-utils"; + +const { useAuthorizedMock } = vi.hoisted(() => ({ useAuthorizedMock: vi.fn() })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: useAuthorizedMock, +})); + +const fetchMock = vi.fn(); + +const requestedUrls = () => fetchMock.mock.calls.map(([url]) => String(url)); + +const renderAs = (userRole: string) => { + useAuthorizedMock.mockReturnValue({ accessToken: "sk-test", userId: "u1", userRole }); + return renderWithProviders(); +}; + +// Deep-linking to /ui/workflows bypasses the sidebar, so the page itself has to +// refuse the render. `/v1/workflows/runs` is proxy-admin-only, so any request +// from a non-admin is the 401 this gate exists to stop. +describe("Workflows page access by role", () => { + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + fetchMock.mockResolvedValue({ + ok: true, + status: 200, + statusText: "OK", + json: async () => ({ runs: [], count: 0 }), + }); + vi.stubGlobal("fetch", fetchMock); + }); + + it("lists workflow runs for an admin", async () => { + renderAs("Admin"); + + await waitFor(() => expect(requestedUrls().some((url) => url.includes("/v1/workflows/runs"))).toBe(true)); + }); + + it.each(["Internal User", "Internal Viewer", "Org Admin", "Unknown Role"])( + "renders the admin-only notice and fires no workflow request for %s", + async (userRole) => { + renderAs(userRole); + + expect(await screen.findByText("Workflow Runs is only available to admin users.")).toBeInTheDocument(); + await waitFor(() => expect(fetchMock).not.toHaveBeenCalled()); + expect(requestedUrls().filter((url) => url.includes("/v1/workflows"))).toEqual([]); + }, + ); + + it("hides the deprecation banner along with the page body for a denied role", () => { + renderAs("Internal User"); + + expect(screen.queryByText(/draft deprecation list/i)).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/workflows/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/workflows/page.tsx index 89dd30f7392..51db7579f82 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/workflows/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/workflows/page.tsx @@ -2,10 +2,18 @@ import WorkflowRuns from "./WorkflowRuns"; import { DeprecationBanner } from "@/components/DeprecationBanner"; +import { AdminOnlyNotice } from "@/components/shared/AdminOnlyNotice"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import useCan from "@/app/(dashboard)/hooks/useCan"; export default function Workflows() { const { accessToken } = useAuthorized(); + const canViewWorkflowRuns = useCan("viewWorkflowRuns"); + + if (!canViewWorkflowRuns) { + return ; + } + return ( <> diff --git a/ui/litellm-dashboard/src/app/globals.css b/ui/litellm-dashboard/src/app/globals.css index 4589d0f528a..97f61a610d9 100644 --- a/ui/litellm-dashboard/src/app/globals.css +++ b/ui/litellm-dashboard/src/app/globals.css @@ -235,3 +235,10 @@ .custom-border { border: 1px solid var(--neutral-border); } + +/* A dialog opened from inside another dialog reads as a drill-down, not a stack: base-ui stamps + this attribute on the parent while a child is open, so the parent steps aside instead of + showing its own edges around a differently sized child. */ +[data-slot="dialog-content"][data-nested-dialog-open] { + visibility: hidden; +} diff --git a/ui/litellm-dashboard/src/autorouter_presets.json b/ui/litellm-dashboard/src/autorouter_presets.json index 58a087d4009..db46da08a3d 100644 --- a/ui/litellm-dashboard/src/autorouter_presets.json +++ b/ui/litellm-dashboard/src/autorouter_presets.json @@ -11,7 +11,8 @@ }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], - "session_affinity": false + "session_affinity": false, + "deployment_affinity": true } }, "openai_family": { @@ -26,7 +27,8 @@ }, "classifier_type": "heuristic", "escalation_keywords": ["LITELLM ESCALATE"], - "session_affinity": false + "session_affinity": false, + "deployment_affinity": true } } } diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index 0f92b61a46e..d0c3235c4e8 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -9,6 +9,7 @@ export interface EntitySpendData { results: any[]; metadata: { total_spend: number; + total_flat_cost?: number; total_api_requests: number; total_successful_requests: number; total_failed_requests: number; @@ -38,6 +39,8 @@ export interface ExportMetadata { export_scope: ExportScope; summary: { total_spend: number; + total_flat_cost?: number; + total_cost?: number; total_requests: number; successful_requests: number; failed_requests: number; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index 2f551176091..08ca298c1f1 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -1861,6 +1861,87 @@ describe("EntityUsageExport utils", () => { expect(result.summary.failed_requests).toBe(20); expect(result.summary.total_tokens).toBe(4500); }); + + it("should include total_flat_cost and total_cost in summary when total_flat_cost is present", () => { + const spendWithFlat: EntitySpendData = { + ...mockSpendData, + metadata: { ...mockSpendData.metadata, total_flat_cost: 6.45 }, + }; + const result = generateMetadata("team", mockDateRange, [], "daily", spendWithFlat); + expect(result.summary.total_flat_cost).toBeCloseTo(6.45, 4); + expect(result.summary.total_cost).toBeCloseTo(46.0 + 6.45, 4); + }); + + it("should omit total_flat_cost and total_cost when total_flat_cost is zero", () => { + const zeroFlat = { ...mockSpendData, metadata: { ...mockSpendData.metadata, total_flat_cost: 0 } }; + const result = generateMetadata("team", mockDateRange, [], "daily", zeroFlat); + expect(result.summary.total_flat_cost).toBeUndefined(); + expect(result.summary.total_cost).toBeUndefined(); + }); + }); + + describe("generateDailyData PTU flat cost", () => { + const dayWithFlat: EntitySpendData = { + results: [ + { + date: "2025-01-01", + breakdown: { + entities: { + "team-1": { + metrics: { + spend: 10, + flat_cost: 6.45, + api_requests: 50, + successful_requests: 50, + failed_requests: 0, + total_tokens: 500, + prompt_tokens: 300, + completion_tokens: 200, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + api_key_breakdown: {}, + }, + }, + }, + }, + ], + metadata: { + total_spend: 10, + total_flat_cost: 6.45, + total_api_requests: 50, + total_successful_requests: 50, + total_failed_requests: 0, + total_tokens: 500, + }, + }; + + it("includes Flat Cost ($) and Total Cost ($) columns when total_flat_cost is present", () => { + const rows = generateDailyData(dayWithFlat, "Team", {}); + expect(rows).toHaveLength(1); + expect(rows[0]).toHaveProperty("Flat Cost ($)"); + expect(rows[0]).toHaveProperty("Total Cost ($)"); + expect(rows[0]["Flat Cost ($)"]).toBe("6.4500"); + expect(rows[0]["Total Cost ($)"]).toBe("16.4500"); + }); + + it("does not include Flat Cost / Total Cost columns when total_flat_cost is zero", () => { + const spendWithoutFlat: EntitySpendData = { + ...dayWithFlat, + metadata: { + total_spend: 10, + total_api_requests: 50, + total_successful_requests: 50, + total_failed_requests: 0, + total_tokens: 500, + total_flat_cost: 0, + }, + }; + const rows = generateDailyData(spendWithoutFlat, "User", {}); + expect(rows).toHaveLength(1); + expect(rows[0]).not.toHaveProperty("Flat Cost ($)"); + expect(rows[0]).not.toHaveProperty("Total Cost ($)"); + }); }); describe("handleExportCSV", () => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 53d100040d7..9adcb50206d 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -110,31 +110,42 @@ export const getEntityBreakdown = ( return Object.values(entitySpend).sort((a, b) => b.metrics.spend - a.metrics.spend); }; +// total_flat_cost defaults to 0 on every entity response, so only a non-zero value +// means a PTU-configured team actually accrued flat cost worth exporting. +const hasFlatCost = (spendData: EntitySpendData): boolean => (spendData.metadata.total_flat_cost ?? 0) > 0; + export const generateDailyData = ( spendData: EntitySpendData, entityLabel: string, teamAliasMap: Record = {}, ): any[] => { const dailyBreakdown: any[] = []; + const includeFlatCost = hasFlatCost(spendData); spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); - dailyBreakdown.push({ + const row: Record = { Date: day.date, [entityLabel]: alias, [`${entityLabel} ID`]: id, "Spend ($)": formatNumberWithCommas(data.metrics.spend, 4), - Requests: data.metrics.api_requests, - "Successful Requests": data.metrics.successful_requests, - "Failed Requests": data.metrics.failed_requests, - "Total Tokens": data.metrics.total_tokens, - "Prompt Tokens": data.metrics.prompt_tokens || 0, - "Completion Tokens": data.metrics.completion_tokens || 0, - "Cache Read Input Tokens": data.metrics.cache_read_input_tokens || 0, - "Cache Creation Input Tokens": data.metrics.cache_creation_input_tokens || 0, - }); + }; + if (includeFlatCost) { + const flatCost = data.metrics.flat_cost || 0; + row["Flat Cost ($)"] = formatNumberWithCommas(flatCost, 4); + row["Total Cost ($)"] = formatNumberWithCommas((data.metrics.spend || 0) + flatCost, 4); + } + row.Requests = data.metrics.api_requests; + row["Successful Requests"] = data.metrics.successful_requests; + row["Failed Requests"] = data.metrics.failed_requests; + row["Total Tokens"] = data.metrics.total_tokens; + row["Prompt Tokens"] = data.metrics.prompt_tokens || 0; + row["Completion Tokens"] = data.metrics.completion_tokens || 0; + row["Cache Read Input Tokens"] = data.metrics.cache_read_input_tokens || 0; + row["Cache Creation Input Tokens"] = data.metrics.cache_creation_input_tokens || 0; + dailyBreakdown.push(row); }); }); @@ -339,23 +350,31 @@ export const generateMetadata = ( selectedFilters: string[], exportScope: ExportScope, spendData: EntitySpendData, -): ExportMetadata => ({ - export_date: new Date().toISOString(), - entity_type: entityType, - date_range: { - from: dateRange.from?.toISOString(), - to: dateRange.to?.toISOString(), - }, - filters_applied: selectedFilters.length > 0 ? selectedFilters : "None", - export_scope: exportScope, - summary: { +): ExportMetadata => { + const summary: ExportMetadata["summary"] = { total_spend: spendData.metadata.total_spend, total_requests: spendData.metadata.total_api_requests, successful_requests: spendData.metadata.total_successful_requests, failed_requests: spendData.metadata.total_failed_requests, total_tokens: spendData.metadata.total_tokens, - }, -}); + }; + if (hasFlatCost(spendData)) { + const flatCost = spendData.metadata.total_flat_cost ?? 0; + summary.total_flat_cost = flatCost; + summary.total_cost = spendData.metadata.total_spend + flatCost; + } + return { + export_date: new Date().toISOString(), + entity_type: entityType, + date_range: { + from: dateRange.from?.toISOString(), + to: dateRange.to?.toISOString(), + }, + filters_applied: selectedFilters.length > 0 ? selectedFilters : "None", + export_scope: exportScope, + summary, + }; +}; export const handleExportCSV = ( spendData: EntitySpendData, diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index e8331294972..c5c77c0c878 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -6,9 +6,14 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import NotificationsManager from "./molecules/notifications_manager"; import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; -import { fetchMCPAccessGroups, getGuardrailsList, teamCreateCall } from "./networking"; +import { fetchMCPAccessGroups, getGuardrailsList, getPoliciesList, teamCreateCall } from "./networking"; import Teams from "./Teams"; +const can = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ + default: (...args: unknown[]) => can(...args), +})); + const mockTeamInfoView = vi.fn(); const mockUseOrganizations = vi.fn(); @@ -173,6 +178,7 @@ const renderWithQueryClient = ( // Re-establish safe defaults before every test (clearAllMocks keeps return values, so restore them here). beforeEach(() => { mockTeamsTableProps = null; + can.mockReturnValue(true); }); describe("Teams - handleCreate organization handling", () => { @@ -956,3 +962,50 @@ describe("Teams - LIT-2530 organization stays optional for proxy admin with a si }); }); }); + +describe("Teams - policies field is gated on the viewPolicies capability", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(getPoliciesList).mockResolvedValue({ policies: [] }); + mockUseOrganizations.mockReturnValue({ data: null }); + }); + + const openAdditionalSettings = async () => { + renderWithQueryClient(); + + act(() => { + fireEvent.click(screen.getAllByRole("button", { name: /create team/i })[0]); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.click(screen.getByText("Additional Settings")); + + await waitFor(() => { + expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); + }); + }; + + it("should render the policies field and load it when the capability is present", async () => { + await openAdditionalSettings(); + + expect(can).toHaveBeenCalledWith("viewPolicies"); + expect(getPoliciesList).toHaveBeenCalledWith("test-token"); + expect(screen.getByText("Policies")).toBeInTheDocument(); + }); + + it("should omit the policies field and skip the admin-only list without the capability", async () => { + can.mockReturnValue(false); + + await openAdditionalSettings(); + + expect(getPoliciesList).not.toHaveBeenCalled(); + expect(screen.queryByText("Policies")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 42ff8bcc4c2..7c269555ec0 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -1,4 +1,5 @@ import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import AvailableTeamsPanel from "@/components/team/AvailableTeamsPanel"; import TeamInfoView from "@/components/team/TeamInfo"; import TeamSSOSettings from "@/components/TeamSSOSettings"; @@ -108,6 +109,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [isTeamDeleting, setIsTeamDeleting] = useState(false); // Add this state near the other useState declarations const [guardrailsList, setGuardrailsList] = useState([]); + const canViewPolicies = useCan("viewPolicies"); const [policiesList, setPoliciesList] = useState([]); const [loggingSettings, setLoggingSettings] = useState([]); const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); @@ -168,8 +170,8 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser }; fetchGuardrails(); - fetchPolicies(); - }, [accessToken]); + if (canViewPolicies) fetchPolicies(); + }, [accessToken, canViewPolicies]); const handleOk = () => { setIsTeamModalVisible(false); @@ -795,36 +797,38 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser } /> - - Policies{" "} - - e.stopPropagation()} - > - - - - - } - name="policies" - className="mt-8" - help="Select existing policies or enter new ones" - > - ({ + value: name, + label: name, + }))} + /> + + )} diff --git a/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index b10fc79be15..420977f8c31 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -1,5 +1,6 @@ export interface SpendMetrics { spend: number; + flat_cost?: number; prompt_tokens: number; completion_tokens: number; total_tokens: number; 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 1e1bb5f9bf4..92dfeab8a7d 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -602,3 +602,32 @@ describe("ComplexityRouterConfig tier labels", () => { expect(screen.getByTitle("Deep")).toBeInTheDocument(); }); }); + +describe("ComplexityRouterConfig affinity panel", () => { + it("holds both affinity switches with their backend defaults", () => { + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Affinity")); + + expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).toBeChecked(); + expect(screen.getByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); + }); + + it("writes deployment_affinity through onChange without touching other keys", () => { + const onChange = vi.fn(); + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Affinity")); + + fireEvent.click(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })); + + expect(onChange).toHaveBeenCalledWith({ ...defaultValue, deployment_affinity: false }); + }); + + it("renders a stored deployment_affinity=false as off", () => { + renderWithProviders( + , + ); + fireEvent.click(screen.getByText("Advanced: Affinity")); + + expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).not.toBeChecked(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index ba805c23c1f..39e6e97bf1d 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -15,6 +15,7 @@ export const DEFAULT_TIER_DISTANCE_PENALTY = 0.5; export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3; export const DEFAULT_CLASSIFIER_CONTEXT_PER_TURN_CHARS = 200; export const DEFAULT_SESSION_AFFINITY = false; +export const DEFAULT_DEPLOYMENT_AFFINITY = true; export interface ComplexityTiers { SIMPLE: string[]; @@ -56,6 +57,7 @@ export interface ComplexityRouterConfigValue { classifier_context_include_assistant_turns?: boolean; classifier_fallback?: ClassifierFallback; session_affinity?: boolean; + deployment_affinity?: boolean; adaptive?: boolean; adaptive_weights?: AdaptiveRouterWeights; tier_distance_penalty?: number; @@ -277,14 +279,26 @@ const ComplexityRouterConfig: React.FC = ({ children: , }, { - key: "session-affinity", + key: "affinity", label: ( - Advanced: Session Affinity + Advanced: Affinity ), children: ( <> +
+ onChange({ ...value, deployment_affinity: deploymentAffinity })} + aria-label="Pin a session to one deployment per model group" + /> + Pin a session to one deployment per model group +
+ + Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off + to load-balance every turn. +
= ({ Pin a session to its first model
- Off by default: every turn is classified on its own merits and routed to the cheapest adequate tier. - Turn this on to reuse the model chosen on a session's first turn for every later turn, which - preserves provider prompt caches and avoids cross-model conversation-history errors, at the cost of - keeping the whole session on the first turn's tier. + Keeps a session on its first turn's model instead of re-classifying each turn. Also pins the + deployment. ), diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index 6a5a1e0f159..7d42d85ccfb 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -273,7 +273,7 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); expandDetailedConfiguration(); - await user.click(screen.getByText("Advanced: Session Affinity")); + await user.click(screen.getByText("Advanced: Affinity")); expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -292,7 +292,7 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); expandDetailedConfiguration(); - await user.click(screen.getByText("Advanced: Session Affinity")); + await user.click(screen.getByText("Advanced: Affinity")); await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); await user.click(screen.getByRole("button", { name: /add auto router/i })); @@ -303,6 +303,46 @@ describe("AddAutoRouterTab", () => { }); }); + it("defaults a new router to deployment affinity on, matching the backend field default", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Affinity")); + expect( + await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }), + ).toBeChecked(); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + deployment_affinity: true, + }); + }); + + it("carries deployment affinity turned off through to the create payload", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders(); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "affinity-router"); + expandDetailedConfiguration(); + await user.click(screen.getByText("Advanced: Affinity")); + await user.click(await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" })); + + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); + expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ + deployment_affinity: false, + }); + }); + // Custom is the escape hatch, not the headline choice, so it's listed after every bundled preset // rather than first. it("lists Custom Configuration after the bundled presets", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 08041413d1f..7b6403f250c 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -15,6 +15,7 @@ import ComplexityRouterConfig, { ComplexityTiers, DEFAULT_ADAPTIVE_WEIGHTS, DEFAULT_SESSION_AFFINITY, + DEFAULT_DEPLOYMENT_AFFINITY, DEFAULT_TIER_DISTANCE_PENALTY, } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; @@ -290,6 +291,7 @@ const AddAutoRouterTab: React.FC = ({ classifierContextIncludeAssistantTurns: complexityRouterConfig.classifier_context_include_assistant_turns, classifierFallback: complexityRouterConfig.classifier_fallback, sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY, + deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx index 9fe36e13998..4e5c5f25374 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx @@ -2,32 +2,37 @@ import { act, fireEvent, render, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import AdvancedSettings from "./advanced_settings"; +const mockUsePtuCostAttributionEnabled = vi.fn(); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled", () => ({ + usePtuCostAttributionEnabled: () => mockUsePtuCostAttributionEnabled(), +})); + +const PTU_LABELS = ["PTU Count", "Calculated Cost per PTU / Hour (USD)", "PTU Effective From (UTC)"]; + +const renderAdvancedSettings = () => + render( + {}} + guardrailsList={[]} + tagsList={{}} + accessToken="test-token" + />, + ); + describe("AdvancedSettings", () => { beforeEach(() => { vi.clearAllMocks(); + mockUsePtuCostAttributionEnabled.mockReturnValue(false); }); + it("should render", () => { - render( - {}} - guardrailsList={[]} - tagsList={{}} - accessToken="test-token" - />, - ); + renderAdvancedSettings(); }); it("should render tags list", async () => { - const { getByText } = render( - {}} - guardrailsList={[]} - tagsList={{}} - accessToken="test-token" - />, - ); + const { getByText } = renderAdvancedSettings(); fireEvent.click(getByText("Advanced Settings")); await waitFor(() => { expect(getByText("Tags")).toBeInTheDocument(); @@ -35,15 +40,7 @@ describe("AdvancedSettings", () => { }); it("should render the litellm params", async () => { - const { getByText } = render( - {}} - guardrailsList={[]} - tagsList={{}} - accessToken="test-token" - />, - ); + const { getByText } = renderAdvancedSettings(); act(() => { fireEvent.click(getByText("Advanced Settings")); }); @@ -51,4 +48,35 @@ describe("AdvancedSettings", () => { expect(getByText("LiteLLM Params")).toBeInTheDocument(); }); }); + + it("hides every PTU field when PTU cost attribution is disabled", async () => { + const { getByText, queryByText } = renderAdvancedSettings(); + act(() => { + fireEvent.click(getByText("Advanced Settings")); + }); + await waitFor(() => { + expect(getByText("Tags")).toBeInTheDocument(); + }); + + for (const label of PTU_LABELS) { + expect(queryByText(label)).not.toBeInTheDocument(); + } + expect(queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); + }); + + it("shows every PTU field when PTU cost attribution is enabled", async () => { + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + const { getByText } = renderAdvancedSettings(); + act(() => { + fireEvent.click(getByText("Advanced Settings")); + }); + + await waitFor(() => { + expect(getByText("PTU Count")).toBeInTheDocument(); + }); + for (const label of PTU_LABELS) { + expect(getByText(label)).toBeInTheDocument(); + } + expect(getByText("PTU Effective To (UTC)")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx index 554fe86bf66..5d7196a84fe 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { Form, Switch, Select, Tooltip } from "antd"; +import { Form, Switch, Select, Tooltip, DatePicker } from "antd"; import { Text, Accordion, AccordionHeader, AccordionBody, TextInput } from "@tremor/react"; import { Row, Col, Typography } from "antd"; import TextArea from "antd/es/input/TextArea"; @@ -9,6 +9,18 @@ import CacheControlSettings from "./cache_control_settings"; import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; import { Tag } from "../tag_management/types"; import { formItemValidateJSON } from "../../utils/textUtils"; +import { + PTU_COUNT_FIELD, + PTU_RATE_FIELD, + PTU_START_FIELD, + ptuCountRules, + ptuPairRule, + ptuRateRules, + ptuStartRequiredRule, + ptuWindowOrderRule, + PTU_END_FIELD, +} from "../../utils/ptuValidation"; +import { usePtuCostAttributionEnabled } from "@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled"; const { Link } = Typography; interface AdvancedSettingsProps { @@ -32,6 +44,7 @@ const AdvancedSettings: React.FC = ({ const [customPricing, setCustomPricing] = React.useState(false); const [pricingModel, setPricingModel] = React.useState<"per_token" | "per_second">("per_token"); const [showCacheControl, setShowCacheControl] = React.useState(false); + const ptuCostAttributionEnabled = usePtuCostAttributionEnabled(); // Add validation function for numbers const validateNumber = (_: any, value: string) => { @@ -182,6 +195,54 @@ const AdvancedSettings: React.FC = ({ />
+ {ptuCostAttributionEnabled && ( + <> + + + + + + + + + + + + + + + + + )} + {customPricing && (
diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 34f7354d412..1a2276f302d 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -26,6 +26,7 @@ const baseParams: BuildComplexityRouterConfigParams = { classifierContextIncludeAssistantTurns: undefined, classifierFallback: undefined, sessionAffinity: false, + deploymentAffinity: true, customTechnicalKeywords: [], keywordTierRules: [], semanticMatchingEnabled: false, @@ -46,6 +47,7 @@ describe("buildComplexityRouterConfig", () => { tiers, classifier_type: "heuristic", session_affinity: false, + deployment_affinity: true, escalation_keywords: ["LITELLM ESCALATE"], }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 9ed1b98db7b..503390aa238 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -30,6 +30,7 @@ export interface BuildComplexityRouterConfigParams { classifierContextIncludeAssistantTurns: boolean | undefined; classifierFallback: ClassifierFallback | undefined; sessionAffinity: boolean; + deploymentAffinity: boolean; customTechnicalKeywords: string[]; keywordTierRules: KeywordTierRule[]; semanticMatchingEnabled: boolean; @@ -53,6 +54,7 @@ export interface ComplexityRouterConfigPayload { classifier_context_include_assistant_turns?: boolean; classifier_fallback?: ClassifierFallback; session_affinity: boolean; + deployment_affinity: boolean; custom_technical_keywords?: string[]; keyword_tier_rules?: { keywords: string[]; tier: KeywordTierRule["tier"] }[]; semantic_keyword_matching?: boolean; @@ -137,6 +139,7 @@ export const buildComplexityRouterConfig = ({ classifierContextIncludeAssistantTurns, classifierFallback, sessionAffinity, + deploymentAffinity, customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, @@ -173,6 +176,7 @@ export const buildComplexityRouterConfig = ({ classifier_context_include_assistant_turns: classifierContextIncludeAssistantTurns, }), session_affinity: sessionAffinity, + deployment_affinity: deploymentAffinity, ...(customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords }), ...(cleanedKeywordTierRules.length > 0 && { keyword_tier_rules: cleanedKeywordTierRules }), escalation_keywords: cleanedEscalationKeywords, diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx index 908a4c498dc..af5fe0931c8 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx @@ -1,6 +1,7 @@ import NotificationManager from "../molecules/notifications_manager"; import { Model, modelCreateCall } from "../networking"; import { provider_map } from "../provider_info_helpers"; +import { ptuPickerToUtcIso } from "../../utils/ptuDatetime"; export const prepareModelAddRequest = async (formValues: Record, accessToken: string, form: any) => { try { @@ -163,6 +164,23 @@ export const prepareModelAddRequest = async (formValues: Record, ac continue; } + // Handle the PTU flat-cost fields (attributed to the team via model_info) + else if (key === "ptu_count" || key === "cost_per_ptu_per_hour") { + if (value !== undefined && value !== null && value !== "") { + modelInfoObj[key] = Number(value); + } + continue; + } + + // Handle the PTU effective window (DatePicker dayjs value -> ISO 8601 UTC string) + else if (key === "ptu_effective_from" || key === "ptu_effective_to") { + const iso = ptuPickerToUtcIso(value); + if (iso !== null) { + modelInfoObj[key] = iso; + } + continue; + } + // Check if key is any of the specified API related keys else { // Add key-value pair to litellm_params dictionary diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 30859b5baea..b56ec734aeb 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -228,6 +228,31 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => { }); }); +describe("buildUpdatedComplexityRouterConfig deployment affinity", () => { + it("writes deployment_affinity=false when the toggle is off", () => { + const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, deployment_affinity: false }); + expect(result.deployment_affinity).toBe(false); + }); + + it("writes deployment_affinity=true when the toggle is on", () => { + const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, deployment_affinity: true }); + expect(result.deployment_affinity).toBe(true); + }); + + it("re-asserts the backend's on-by-default when the form value is absent, rather than dropping the key", () => { + const result = buildUpdatedComplexityRouterConfig({ ...STORED, deployment_affinity: false }, FORM_VALUE); + expect(result.deployment_affinity).toBe(true); + }); + + it("stops a stored deployment_affinity=false from surviving a save that turned the toggle back on", () => { + const result = buildUpdatedComplexityRouterConfig( + { ...STORED, deployment_affinity: false }, + { ...FORM_VALUE, deployment_affinity: true }, + ); + expect(result.deployment_affinity).toBe(true); + }); +}); + describe("buildUpdatedComplexityRouterConfig tier labels", () => { const RENAMED = { ...STORED, tier_labels: { SIMPLE: "Cheap", REASONING: "Deep" } }; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts index eb5bb46f0e8..027e01a9351 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.ts @@ -48,6 +48,7 @@ const expectedClassifiedTierConfig = { embedding_model: "voyage-4-large", match_threshold: 0.65, session_affinity: false, + deployment_affinity: true, adaptive: true, adaptive_weights: { quality: 0.4, cost: 0.6 }, adaptive_eligible: "classified_tier", @@ -68,6 +69,7 @@ const expectedAdaptiveDisabledConfig = { embedding_model: "voyage-4-large", match_threshold: 0.65, session_affinity: false, + deployment_affinity: true, }; describe("buildUpdatedComplexityRouterConfig", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 9b36218e17a..74dea8cc2ed 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -6,12 +6,19 @@ import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../te import NotificationsManager from "@/components/molecules/notifications_manager"; import EditAutoRouterModal from "./edit_auto_router_modal"; -const { modelPatchUpdateCall, modelAvailableCall } = vi.hoisted(() => ({ +const { modelPatchUpdateCall, modelAvailableCall, getAutoRouterClassifierDefaultPromptCall } = vi.hoisted(() => ({ modelPatchUpdateCall: vi.fn().mockResolvedValue({}), modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), + getAutoRouterClassifierDefaultPromptCall: vi.fn().mockResolvedValue("Classify the request into exactly one tier."), })); -vi.mock("../networking", () => ({ modelPatchUpdateCall, modelAvailableCall })); +vi.mock("../networking", () => ({ + modelPatchUpdateCall, + modelAvailableCall, + getAutoRouterClassifierDefaultPromptCall, +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "sk-test" }) })); vi.mock("@/components/llm_calls/fetch_models", () => ({ fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "gpt-4o-mini" }]), @@ -243,6 +250,22 @@ describe("EditAutoRouterModal classifier context window", () => { expect(config.classifier_context_per_turn_chars).toBe(300); }); + // The prompt editor is a base-ui Dialog at z-index 50. Housing this form in an antd Modal put a + // z-index 1000 overlay between the operator and it, so the editor opened underneath and could + // not be read or typed into. jsdom does not paint, so the assertion is the invariant behind the + // stacking: both overlays come from the one Dialog primitive the create form already uses. + it("opens the classifier prompt editor in the same overlay layer as the form", async () => { + const user = userEvent.setup(); + const { baseElement } = renderLlmModal(); + + await user.click(await screen.findByText("Advanced: Classification Method")); + await user.click(await screen.findByRole("button", { name: /prompt/i })); + + expect(await screen.findByLabelText("Classifier system prompt")).toBeInTheDocument(); + expect(baseElement.querySelectorAll('[data-slot="dialog-content"]')).toHaveLength(2); + expect(baseElement.querySelector(".ant-modal")).toBeNull(); + }); + it("persists an edited classifier context window size", async () => { const user = userEvent.setup(); renderLlmModal(); @@ -342,7 +365,7 @@ describe("EditAutoRouterModal session affinity", () => { const user = userEvent.setup(); renderWithStoredConfig(STORED_CONFIG); - await user.click(await screen.findByText("Advanced: Session Affinity")); + await user.click(await screen.findByText("Advanced: Affinity")); expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).not.toBeChecked(); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -355,7 +378,7 @@ describe("EditAutoRouterModal session affinity", () => { const user = userEvent.setup(); renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); - await user.click(await screen.findByText("Advanced: Session Affinity")); + await user.click(await screen.findByText("Advanced: Affinity")); expect(await screen.findByRole("switch", { name: "Pin a session to its first model" })).toBeChecked(); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -368,7 +391,7 @@ describe("EditAutoRouterModal session affinity", () => { const user = userEvent.setup(); renderWithStoredConfig(STORED_CONFIG); - await user.click(await screen.findByText("Advanced: Session Affinity")); + await user.click(await screen.findByText("Advanced: Affinity")); await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -381,7 +404,7 @@ describe("EditAutoRouterModal session affinity", () => { const user = userEvent.setup(); renderWithStoredConfig({ ...STORED_CONFIG, session_affinity: true }); - await user.click(await screen.findByText("Advanced: Session Affinity")); + await user.click(await screen.findByText("Advanced: Affinity")); await user.click(await screen.findByRole("switch", { name: "Pin a session to its first model" })); await user.click(screen.getByRole("button", { name: /save changes/i })); @@ -391,6 +414,67 @@ describe("EditAutoRouterModal session affinity", () => { }); }); +describe("EditAutoRouterModal deployment affinity", () => { + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + const renderWithStoredConfig = (complexity_router_config: Record) => + renderWithProviders( + , + ); + + it("shows a stored config with no deployment_affinity key as on, matching the backend default", async () => { + const user = userEvent.setup(); + renderWithStoredConfig(STORED_CONFIG); + + await user.click(await screen.findByText("Advanced: Affinity")); + expect( + await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }), + ).toBeChecked(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().deployment_affinity).toBe(true); + }); + + it("shows a stored deployment_affinity=false as off and preserves it through an untouched save", async () => { + const user = userEvent.setup(); + renderWithStoredConfig({ ...STORED_CONFIG, deployment_affinity: false }); + + await user.click(await screen.findByText("Advanced: Affinity")); + expect( + await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" }), + ).not.toBeChecked(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().deployment_affinity).toBe(false); + }); + + it("persists turning deployment affinity off", async () => { + const user = userEvent.setup(); + renderWithStoredConfig(STORED_CONFIG); + + await user.click(await screen.findByText("Advanced: Affinity")); + await user.click(await screen.findByRole("switch", { name: "Pin a session to one deployment per model group" })); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled()); + expect(savedConfig().deployment_affinity).toBe(false); + }); +}); + describe("EditAutoRouterModal custom classifier prompt and fallback", () => { beforeEach(() => { modelPatchUpdateCall.mockClear(); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index cf1a5727948..6818d3b850c 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -1,6 +1,6 @@ import React, { useEffect, useState } from "react"; -import { Modal, Form, Button, Select as AntdSelect, Tooltip } from "antd"; -import { Text, TextInput } from "@tremor/react"; +import { Form, Button, Select as AntdSelect, Tooltip } from "antd"; +import { TextInput } from "@tremor/react"; import { modelAvailableCall, modelPatchUpdateCall } from "../networking"; import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder from "../add_model/RouterConfigBuilder"; @@ -21,9 +21,18 @@ import ComplexityRouterConfig, { ComplexityRouterConfigValue, DEFAULT_ADAPTIVE_WEIGHTS, DEFAULT_SESSION_AFFINITY, + DEFAULT_DEPLOYMENT_AFFINITY, DEFAULT_TIER_DISTANCE_PENALTY, } from "../add_model/ComplexityRouterConfig"; import NotificationsManager from "../molecules/notifications_manager"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; interface EditAutoRouterModalProps { isVisible: boolean; @@ -47,6 +56,7 @@ const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "classifier_context_include_assistant_turns", "classifier_fallback", "session_affinity", + "deployment_affinity", "adaptive", "adaptive_weights", "tier_distance_penalty", @@ -119,6 +129,7 @@ export const buildUpdatedComplexityRouterConfig = ( classifier_context_include_assistant_turns: value.classifier_context_include_assistant_turns, }), session_affinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, + deployment_affinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, ...(customTechnicalKeywords && customTechnicalKeywords.length > 0 && { custom_technical_keywords: customTechnicalKeywords, @@ -255,6 +266,10 @@ const EditAutoRouterModal: React.FC = ({ typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, + deployment_affinity: + typeof parsedConfig.deployment_affinity === "boolean" + ? parsedConfig.deployment_affinity + : DEFAULT_DEPLOYMENT_AFFINITY, adaptive: parsedConfig.adaptive || false, adaptive_weights: parsedConfig.adaptive_weights, tier_distance_penalty: parsedConfig.tier_distance_penalty, @@ -432,27 +447,14 @@ const EditAutoRouterModal: React.FC = ({ })); return ( - - Cancel - , - - - , - ]} - width={1000} - destroyOnHidden - > -
- - Edit the auto router configuration including routing logic, default models, and access settings. - + !open && onCancel()}> + + + Edit Auto Router Configuration + + Edit the auto router configuration including routing logic, default models, and access settings. + +
{/* Auto Router Name */} @@ -552,8 +554,17 @@ const EditAutoRouterModal: React.FC = ({ )} -
-
+ + + + + + + + + ); }; diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts index 637325ad98d..15c45153026 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.test.ts @@ -1,11 +1,12 @@ -import { describe, expect, it, vi } from "vitest"; -import { fetchTeamFilterOptions } from "./filter_helpers"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { fetchAllTeams, fetchTeamFilterOptions } from "./filter_helpers"; const mockKeyListCall = vi.fn(); +const mockTeamListCall = vi.fn(); vi.mock("@/components/networking", () => ({ keyListCall: (...args: unknown[]) => mockKeyListCall(...args), - teamListCall: vi.fn(), + teamListCall: (...args: unknown[]) => mockTeamListCall(...args), organizationListCall: vi.fn(), })); @@ -78,3 +79,39 @@ describe("fetchTeamFilterOptions", () => { expect(result).toEqual({ keyAliases: [], organizationIds: [], userIds: [] }); }); }); + +describe("fetchAllTeams", () => { + beforeEach(() => { + mockTeamListCall.mockReset(); + }); + + it("forwards the scoping user id to /team/list and returns the rows it answers with", async () => { + mockTeamListCall.mockResolvedValue([{ team_id: "team-a" }, { team_id: "team-b" }]); + + const teams = await fetchAllTeams("tok-123", null, "member-7"); + + expect(mockTeamListCall).toHaveBeenCalledWith("tok-123", null, "member-7"); + expect(teams.map((team) => team.team_id)).toEqual(["team-a", "team-b"]); + }); + + it("sends no user id when the caller is entitled to the broad list", async () => { + mockTeamListCall.mockResolvedValue([]); + + await fetchAllTeams("tok-123"); + + expect(mockTeamListCall).toHaveBeenCalledWith("tok-123", null, null); + }); + + it("keeps the organization filter independent of the scoping user id", async () => { + mockTeamListCall.mockResolvedValue([]); + + await fetchAllTeams("tok-123", "org-1", "member-7"); + + expect(mockTeamListCall).toHaveBeenCalledWith("tok-123", "org-1", "member-7"); + }); + + it("returns an empty list without calling the endpoint when there is no access token", async () => { + expect(await fetchAllTeams(null, null, "member-7")).toEqual([]); + expect(mockTeamListCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts index fb701b4656b..7eef4d3a8b3 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_helpers.ts @@ -114,9 +114,15 @@ export const fetchTeamFilterOptions = async ( * Fetches all teams across all pages * @param accessToken The access token for API authentication * @param organizationId Optional organization ID to filter teams + * @param userID Scopes the list to that user's teams. Required for roles the endpoint + * does not grant a broad list to; see `teamListScopeUserId` * @returns Array of all teams */ -export const fetchAllTeams = async (accessToken: string | null, organizationId?: string | null): Promise => { +export const fetchAllTeams = async ( + accessToken: string | null, + organizationId?: string | null, + userID?: string | null, +): Promise => { if (!accessToken) return []; try { @@ -125,7 +131,7 @@ export const fetchAllTeams = async (accessToken: string | null, organizationId?: let hasMorePages = true; while (hasMorePages) { - const response = await teamListCall(accessToken, organizationId || null, null); + const response = await teamListCall(accessToken, organizationId || null, userID ?? null); // Add teams from this page allTeams = [...allTeams, ...response]; diff --git a/ui/litellm-dashboard/src/components/leftnav.test.tsx b/ui/litellm-dashboard/src/components/leftnav.test.tsx index a5b273a0f56..8eca990261c 100644 --- a/ui/litellm-dashboard/src/components/leftnav.test.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.test.tsx @@ -3,9 +3,12 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../tests/test-utils"; import Sidebar, { menuGroups, getBreadcrumb } from "./leftnav"; -vi.mock("../utils/roles", () => { +vi.mock("../utils/roles", async (importOriginal) => { + const actual = await importOriginal(); return { + ...actual, all_admin_roles: ["admin", "admin_viewer"], + old_admin_roles: ["admin", "admin_viewer"], internalUserRoles: ["internal"], rolesWithWriteAccess: ["admin", "internal"], rolesAllowedToViewWriteScopedPages: ["admin", "internal", "admin_viewer"], @@ -91,6 +94,11 @@ describe("Sidebar (leftnav)", () => { collapsed: false, }; + afterEach(() => { + mockUseAuthorized.mockReset(); + mockUseOrganizations.mockReset(); + }); + it("should link the logo to the UI home route rather than the proxy origin", () => { renderWithProviders(); @@ -174,19 +182,19 @@ describe("Sidebar (leftnav)", () => { }; it("hides Playground from Admin Viewer (cost-incurring action)", () => { - mockUseAuthorized.mockReturnValueOnce(adminViewerAuth); + mockUseAuthorized.mockReturnValue(adminViewerAuth); renderWithProviders(); expect(screen.queryByText("Playground")).not.toBeInTheDocument(); }); it("shows Models + Endpoints to Admin Viewer (read-only)", () => { - mockUseAuthorized.mockReturnValueOnce(adminViewerAuth); + mockUseAuthorized.mockReturnValue(adminViewerAuth); renderWithProviders(); expect(screen.getByText("Models + Endpoints")).toBeInTheDocument(); }); it("shows Agents (under Agentic) to Admin Viewer (read-only)", async () => { - mockUseAuthorized.mockReturnValueOnce(adminViewerAuth); + mockUseAuthorized.mockReturnValue(adminViewerAuth); renderWithProviders(); // Agents is now nested under the "Agentic" submenu — expand parent // first to render the children, then assert Agents is visible. @@ -199,7 +207,7 @@ describe("Sidebar (leftnav)", () => { }); it("shows Logs to Admin Viewer", () => { - mockUseAuthorized.mockReturnValueOnce(adminViewerAuth); + mockUseAuthorized.mockReturnValue(adminViewerAuth); renderWithProviders(); expect(screen.getByText("Logs")).toBeInTheDocument(); }); @@ -210,6 +218,7 @@ describe("Sidebar (leftnav)", () => { userId: "internal-user-id", accessToken: "test-access-token", userRole: "internal", + isViewOnly: false, token: "test-token", userEmail: "internal@example.com", premiumUser: false, @@ -244,13 +253,149 @@ describe("Sidebar (leftnav)", () => { expect(screen.getByText("Tool Policies")).toBeInTheDocument(); }); }); + + it("should hide the Policies entry from internal users while keeping Guardrails", () => { + mockUseAuthorized.mockReturnValue(internalAuth); + renderWithProviders(); + + expect(screen.getByText("Guardrails")).toBeInTheDocument(); + expect(screen.queryByText("Policies")).not.toBeInTheDocument(); + }); + + it("should hide the Prompts entry from internal users while keeping other Experimental children", async () => { + mockUseAuthorized.mockReturnValue(internalAuth); + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByText("Experimental")); + }); + await waitFor(() => { + expect(screen.getByText("API Playground")).toBeInTheDocument(); + }); + expect(screen.queryByText("Prompts")).not.toBeInTheDocument(); + }); + + it("should hide Old Usage from internal users while keeping other Experimental children", async () => { + mockUseAuthorized.mockReturnValue(internalAuth); + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByText("Experimental")); + }); + await waitFor(() => { + expect(screen.getByText("API Playground")).toBeInTheDocument(); + }); + expect(screen.queryByText("Old Usage")).not.toBeInTheDocument(); + }); + + it("should show Old Usage to admins", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByText("Experimental")); + }); + await waitFor(() => { + expect(screen.getByText("Old Usage")).toBeInTheDocument(); + }); + }); + }); + + // Workflow Runs, Memory and Guardrails Monitor render a shell and then 401 + // for every non-proxy-admin role, because their page-load routes sit outside + // internal_user_routes / self_managed_routes. Cost Optimization does not: + // its primary call is /user/daily/activity, which every role may make, so + // the entry stays and only its proxy-wide tabs are gated inside the page. + describe("capability-gated pages whose data is proxy-admin-only", () => { + const authFor = (userRole: string) => ({ + userId: "some-user-id", + accessToken: "test-access-token", + userRole, + isViewOnly: false, + token: "test-token", + userEmail: "someone@example.com", + premiumUser: false, + disabledPersonalKeyCreation: false, + showSSOBanner: false, + }); + + afterEach(() => { + mockUseAuthorized.mockReset(); + }); + + it("hides Workflow Runs and Memory from an internal user under Agentic", async () => { + mockUseAuthorized.mockReturnValue(authFor("internal")); + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByText("Agentic")); + }); + // Liveness gate: the sibling Agents child stays visible to this role, so + // the absences below mean the gate fired, not that the group never opened. + await waitFor(() => { + expect(screen.getByText("Agents")).toBeInTheDocument(); + }); + expect(screen.queryByText("Workflow Runs")).not.toBeInTheDocument(); + expect(screen.queryByText("Memory")).not.toBeInTheDocument(); + }); + + // An org admin's session role is "Org Admin", which no capability list + // carries, and the proxy denies these routes to org admins too because + // `_user_is_org_admin` needs an organization_id the page-load GET never sends. + // Agents is already out of reach for this role, so gating the other two + // empties the Agentic group entirely and the parent must go with it rather + // than degrade into a leaf link to the non-route `?page=agentic`. + it("drops the whole Agentic group for an org admin once its last child is gated", () => { + mockUseAuthorized.mockReturnValue(authFor("org_admin")); + renderWithProviders(); + + // Liveness gate: Logs carries no role list, so it proves the sidebar rendered. + expect(screen.getByText("Logs")).toBeInTheDocument(); + expect(screen.queryByText("Agentic")).not.toBeInTheDocument(); + expect(screen.queryByText("Workflow Runs")).not.toBeInTheDocument(); + expect(screen.queryByText("Memory")).not.toBeInTheDocument(); + }); + + it("keeps the Agentic group for an internal user, who can still see Agents", () => { + mockUseAuthorized.mockReturnValue(authFor("internal")); + renderWithProviders(); + + expect(screen.getByText("Agentic")).toBeInTheDocument(); + }); + + it("shows Workflow Runs and Memory to admins", async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByText("Agentic")); + }); + await waitFor(() => { + expect(screen.getByText("Workflow Runs")).toBeInTheDocument(); + }); + expect(screen.getByText("Memory")).toBeInTheDocument(); + }); + + it("hides Guardrails Monitor from an internal user while keeping Usage and Cost Optimization", () => { + mockUseAuthorized.mockReturnValue(authFor("internal")); + renderWithProviders(); + + expect(screen.queryByText("Guardrails Monitor")).not.toBeInTheDocument(); + expect(screen.getByText("Usage")).toBeInTheDocument(); + expect(screen.getByText("Cost Optimization")).toBeInTheDocument(); + }); + + it("shows Guardrails Monitor to admins", () => { + renderWithProviders(); + + expect(screen.getByText("Guardrails Monitor")).toBeInTheDocument(); + }); }); it("should show Organizations tab for organization admins", () => { - mockUseAuthorized.mockReturnValueOnce({ + mockUseAuthorized.mockReturnValue({ userId: "org-admin-user-id", accessToken: "test-access-token", userRole: "viewer", + isViewOnly: false, token: "test-token", userEmail: "orgadmin@example.com", premiumUser: false, @@ -258,7 +403,7 @@ describe("Sidebar (leftnav)", () => { showSSOBanner: false, }); - mockUseOrganizations.mockReturnValueOnce({ + mockUseOrganizations.mockReturnValue({ data: [ { organization_id: "org-1", diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index f08092d0e38..2ece64271f8 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -1,6 +1,6 @@ -import { useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import useIsOrgAdmin from "@/app/(dashboard)/hooks/useIsOrgAdmin"; import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails"; import { useLogout } from "@/app/(dashboard)/hooks/useLogout"; import { getProxyBaseUrl } from "@/components/networking"; @@ -75,7 +75,6 @@ import { } from "../utils/roles"; import BetaBadge from "./BetaBadge"; import NewBadge from "./common_components/NewBadge"; -import type { Organization } from "./networking"; import SidebarAccountMenu from "./SidebarAccountMenu/SidebarAccountMenu"; import SidebarUsageCard from "./SidebarUsageCard"; import { MIGRATED_PAGES, migratedHref, legacyPageHref } from "@/utils/migratedPages"; @@ -146,8 +145,20 @@ const menuGroups: MenuGroup[] = [ icon: , roles: rolesAllowedToViewWriteScopedPages, }, - { key: "workflows", page: "workflows", label: "Workflow Runs", icon: }, - { key: "memory", page: "memory", label: "Memory", icon: }, + { + key: "workflows", + page: "workflows", + label: "Workflow Runs", + icon: , + roles: rolesWithCapability("viewWorkflowRuns"), + }, + { + key: "memory", + page: "memory", + label: "Memory", + icon: , + roles: rolesWithCapability("viewMemory"), + }, ], }, { key: "mcp-servers", page: "mcp-servers", label: "MCP Servers", icon: }, @@ -158,7 +169,7 @@ const menuGroups: MenuGroup[] = [ page: "policies", label: "Policies", icon: , - roles: all_admin_roles, + roles: rolesWithCapability("viewPolicies"), }, { key: "tools", @@ -206,7 +217,7 @@ const menuGroups: MenuGroup[] = [ page: "guardrails-monitor", label: "Guardrails Monitor", icon: , - roles: [...all_admin_roles, ...internalUserRoles], + roles: rolesWithCapability("viewGuardrailUsage"), }, ], }, @@ -268,7 +279,13 @@ const menuGroups: MenuGroup[] = [ label: "Experimental", icon: , children: [ - { key: "prompts", page: "prompts", label: "Prompts", icon: , roles: all_admin_roles }, + { + key: "prompts", + page: "prompts", + label: "Prompts", + icon: , + roles: rolesWithCapability("viewPrompts"), + }, { key: "transform-request", page: "transform-request", @@ -283,7 +300,13 @@ const menuGroups: MenuGroup[] = [ icon: , roles: all_admin_roles, }, - { key: "4", page: "usage", label: "Old Usage", icon: }, + { + key: "4", + page: "usage", + label: "Old Usage", + icon: , + roles: rolesWithCapability("viewGlobalSpend"), + }, ], }, ], @@ -408,7 +431,7 @@ const Sidebar_: React.FC = ({ allowVectorStoresForTeamAdmins, }) => { const { userId, accessToken, userRole, isViewOnly } = useAuthorized(); - const { data: organizations } = useOrganizations(); + const isOrgAdmin = useIsOrgAdmin(); const { data: teams } = useTeams(); const { logoUrl } = useTheme(); const { data: healthData } = useHealthReadinessDetails(accessToken); @@ -435,13 +458,6 @@ const Sidebar_: React.FC = ({ } } - const isOrgAdmin = useMemo(() => { - if (!userId || !organizations) return false; - return organizations.some((org: Organization) => - org.members?.some((member) => member.user_id === userId && member.user_role === "org_admin"), - ); - }, [userId, organizations]); - const isTeamAdmin = useMemo(() => isUserTeamAdminForAnyTeam(teams ?? null, userId ?? ""), [teams, userId]); const filterItemsByRole = (items: MenuItem[]): MenuItem[] => { @@ -449,6 +465,9 @@ const Sidebar_: React.FC = ({ return items .map((item) => ({ ...item, children: item.children ? filterItemsByRole(item.children) : undefined })) .filter((item) => { + // A parent whose children were all filtered out renders as a leaf link + // to its own page id, which is not a real route. Drop it instead. + if (item.children && item.children.length === 0) return false; if (item.key === "llm-playground" && isViewOnly) return false; if (item.key === "organizations" || item.key === "users") { const hasRoleAccess = !item.roles || item.roles.includes(userRole) || isOrgAdmin; diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 326b49ff896..a3ce40494eb 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -47,6 +47,11 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ useModelCostMap: (...args: any[]) => mockUseModelCostMap(...args), })); +const mockUsePtuCostAttributionEnabled = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled", () => ({ + usePtuCostAttributionEnabled: () => mockUsePtuCostAttributionEnabled(), +})); + const mockNotificationsManager = vi.mocked(NotificationsManager); const mockModelInfoV1Call = vi.mocked(networking.modelInfoV1Call); const mockCredentialGetCall = vi.mocked(networking.credentialGetCall); @@ -99,6 +104,7 @@ describe("ModelInfoView", () => { }, }); vi.clearAllMocks(); + mockUsePtuCostAttributionEnabled.mockReturnValue(false); mockUseModelsInfo.mockReturnValue({ data: { @@ -608,6 +614,100 @@ describe("ModelInfoView", () => { expect(updatePayload.litellm_params).not.toHaveProperty("vector_store_ids"); }); + describe("PTU cost attribution gate", () => { + const ptuModelData = { + ...defaultModelData, + model_info: { + ...defaultModelData.model_info, + team_id: "team-1", + ptu_count: 15, + cost_per_ptu_per_hour: 2, + ptu_effective_from: "2026-07-01T00:00:00+00:00", + ptu_effective_to: "2026-08-01T00:00:00+00:00", + }, + }; + + const renderWithPtuModel = () => { + mockUseModelsInfo.mockReturnValue({ data: { data: [ptuModelData] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [ptuModelData] }); + return render(, { wrapper }); + }; + + it("hides the PTU fields when disabled, even for a model that already stores PTU config", async () => { + renderWithPtuModel(); + + await waitFor(() => { + expect(screen.getByText("Model Settings")).toBeInTheDocument(); + }); + + expect(screen.queryByText("PTU Count")).not.toBeInTheDocument(); + expect(screen.queryByText("Cost per PTU / Hour (USD)")).not.toBeInTheDocument(); + expect(screen.queryByText("PTU Effective From (UTC)")).not.toBeInTheDocument(); + expect(screen.queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); + }); + + it("shows the PTU fields when enabled", async () => { + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + renderWithPtuModel(); + + await waitFor(() => { + expect(screen.getByText("PTU Count")).toBeInTheDocument(); + }); + expect(screen.getByText("Cost per PTU / Hour (USD)")).toBeInTheDocument(); + expect(screen.getByText("PTU Effective From (UTC)")).toBeInTheDocument(); + expect(screen.getByText("PTU Effective To (UTC)")).toBeInTheDocument(); + }); + + it("omits PTU fields from the save payload when disabled, so an unrelated edit cannot clear stored config", async () => { + const user = userEvent.setup(); + renderWithPtuModel(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(mockModelPatchUpdateCall).toHaveBeenCalled(); + }); + + const modelInfo = mockModelPatchUpdateCall.mock.calls[0][1].model_info; + expect(modelInfo).not.toHaveProperty("ptu_count"); + expect(modelInfo).not.toHaveProperty("cost_per_ptu_per_hour"); + expect(modelInfo).not.toHaveProperty("ptu_effective_from"); + expect(modelInfo).not.toHaveProperty("ptu_effective_to"); + }); + + it("sends the PTU fields on save when enabled", async () => { + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + const user = userEvent.setup(); + renderWithPtuModel(); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /edit settings/i })); + + await waitFor(() => { + expect(screen.getByRole("button", { name: /save changes/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(mockModelPatchUpdateCall).toHaveBeenCalled(); + }); + + const modelInfo = mockModelPatchUpdateCall.mock.calls[0][1].model_info; + expect(modelInfo.ptu_count).toBe(15); + expect(modelInfo.cost_per_ptu_per_hour).toBe(2); + }); + }); + it("should not include input_cost_per_token or output_cost_per_token in update payload when user does not touch cost fields", async () => { // Regression: editing a model without touching cost fields used to inject // input_cost_per_token: 0 and output_cost_per_token: 0 into litellm_params, diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index c3327911993..e1a4d4311ab 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -17,7 +17,21 @@ import { Title, Button as TremorButton, } from "@tremor/react"; -import { Button, Form, Input, Modal, Select, Tooltip } from "antd"; +import { Button, DatePicker, Form, Input, Modal, Select, Tooltip } from "antd"; +import { formatPtuUtcDisplay, utcIsoToPickerValue } from "../utils/ptuDatetime"; +import { applyPtuModelInfo } from "../utils/ptuModelInfo"; +import { usePtuCostAttributionEnabled } from "@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled"; +import { + PTU_COUNT_FIELD, + PTU_RATE_FIELD, + ptuCountRules, + ptuPairRule, + ptuRateRules, + ptuStartRequiredRule, + ptuWindowOrderRule, + PTU_END_FIELD, + PTU_START_FIELD, +} from "../utils/ptuValidation"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import { CheckIcon, CopyIcon } from "lucide-react"; import { useEffect, useMemo, useState } from "react"; @@ -67,6 +81,62 @@ interface ModelInfoViewProps { modelAccessGroups: string[] | null; } +interface PtuEditField { + name: string; + label: string; + input: "number" | "datetime"; + placeholder?: string; + isCount?: boolean; + isRate?: boolean; + isStart?: boolean; + pairedWith?: string; + windowPeer?: string; + bound?: "start" | "end"; +} + +const PTU_EDIT_FIELDS: PtuEditField[] = [ + { + name: PTU_COUNT_FIELD, + label: "PTU Count", + input: "number", + placeholder: "e.g. 15", + isCount: true, + pairedWith: PTU_RATE_FIELD, + }, + { + name: PTU_RATE_FIELD, + label: "Cost per PTU / Hour (USD)", + input: "number", + placeholder: "e.g. 2.00", + isRate: true, + pairedWith: PTU_COUNT_FIELD, + }, + { + name: PTU_START_FIELD, + label: "PTU Effective From (UTC)", + input: "datetime", + isStart: true, + windowPeer: PTU_END_FIELD, + bound: "start", + }, + { + name: PTU_END_FIELD, + label: "PTU Effective To (UTC)", + input: "datetime", + windowPeer: PTU_START_FIELD, + bound: "end", + }, +]; + +const ptuFieldDependencies = ({ isStart, pairedWith, windowPeer }: PtuEditField): string[] | undefined => { + const deps = [ + ...(isStart ? [PTU_COUNT_FIELD] : []), + ...(pairedWith ? [pairedWith] : []), + ...(windowPeer ? [windowPeer] : []), + ]; + return deps.length ? deps : undefined; +}; + interface ComplexityRouterTierConfig { tiers?: { SIMPLE?: unknown; @@ -156,6 +226,7 @@ export default function ModelInfoView({ const { data: modelCostMapData } = useModelCostMap(); const { data: modelHubData } = useModelHub(); const { data: teams } = useTeams(); + const ptuCostAttributionEnabled = usePtuCostAttributionEnabled(); // Transform the model data const getProviderFromModel = (model: string) => { @@ -427,6 +498,7 @@ export default function ModelInfoView({ health_check_model: values.health_check_model, }; } + updatedModelInfo = applyPtuModelInfo(updatedModelInfo, values, ptuCostAttributionEnabled); } catch (e) { NotificationsManager.fromBackend("Invalid JSON in Model Info"); return; @@ -769,6 +841,10 @@ export default function ModelInfoView({ output_cost: localModelData.litellm_params?.output_cost_per_token ? localModelData.litellm_params.output_cost_per_token * 1_000_000 : localModelData.model_info?.output_cost_per_token * 1_000_000 || null, + ptu_count: localModelData.model_info?.ptu_count ?? null, + cost_per_ptu_per_hour: localModelData.model_info?.cost_per_ptu_per_hour ?? null, + ptu_effective_from: utcIsoToPickerValue(localModelData.model_info?.ptu_effective_from), + ptu_effective_to: utcIsoToPickerValue(localModelData.model_info?.ptu_effective_to), cache_read_cost: localModelData.litellm_params?.cache_read_input_token_cost !== undefined && localModelData.litellm_params?.cache_read_input_token_cost !== null @@ -872,6 +948,47 @@ export default function ModelInfoView({ )}
+ {ptuCostAttributionEnabled && + PTU_EDIT_FIELDS.map((ptuField) => { + const { name, label, input, placeholder, isCount, isRate, isStart, pairedWith } = ptuField; + const { windowPeer, bound } = ptuField; + return ( +
+ {label} + {isEditing ? ( + + {input === "number" ? ( + + ) : ( + + )} + + ) : ( +
+ {(input === "datetime" + ? formatPtuUtcDisplay(localModelData?.model_info?.[name]) + : localModelData?.model_info?.[name]) ?? "Not Set"} +
+ )} +
+ ); + })} +
Cache Read Cost (per 1M tokens) {isEditing ? ( diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 17a5ca37990..321dfd6b2a6 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -66,6 +66,7 @@ import type { } from "@/app/(dashboard)/caching/_components/coordination_redis_settings/types"; import { MCP_TOOLS_PREVIEW_FORBIDDEN_MESSAGE } from "./mcp_tools/constants"; import type { ComplexityRouterConfigPayload } from "./add_model/build_complexity_router_config"; +import type { VectorStoreIndex } from "@/app/(dashboard)/vector-stores/_components/IndexesTab"; import type { RoutingDecision } from "./view_logs/LogDetailsDrawer/RoutingDecisionCard"; import { createApiClient, @@ -2098,7 +2099,6 @@ export const adminTopEndUsersCall = async ( export const adminspendByProvider = async ( accessToken: string, - keyToken: string | null, startTime: string | undefined, endTime: string | undefined, ) => { @@ -2107,7 +2107,6 @@ export const adminspendByProvider = async ( accessToken, query: { ...(startTime && endTime ? { start_date: startTime, end_date: endTime } : {}), - ...(keyToken ? { api_key: keyToken } : {}), }, }); return data; @@ -5622,6 +5621,20 @@ export const vectorStoreListCall = async ( } }; +export interface IndexesListResponse { + object: string; + data: VectorStoreIndex[]; +} + +export const indexesListCall = async (accessToken: string): Promise => { + try { + return await apiClient.get(`/v1/indexes`, { accessToken }); + } catch (error) { + console.error("Error listing indexes:", error); + throw error; + } +}; + export const vectorStoreDeleteCall = async (accessToken: string, vectorStoreId: string): Promise => { try { let url = proxyBaseUrl ? `${proxyBaseUrl}/vector_store/delete` : `/vector_store/delete`; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx index 2f72fa37d86..5071fe6ff95 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.test.tsx @@ -2,7 +2,7 @@ import { act, fireEvent, within } from "@testing-library/react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, screen, waitFor } from "../../../tests/test-utils"; import { Team } from "../key_team_helpers/key_list"; -import { userFilterUICall } from "../networking"; +import { getPoliciesList, getPromptsList, userFilterUICall } from "../networking"; import CreateKey from "./create_key_button"; const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall, teamDropdownTeamsRef } = @@ -777,4 +777,47 @@ describe("CreateKey", () => { }); }); }); + + describe("policy and prompt fields", () => { + const POLICIES_PLACEHOLDER = "Premium feature - Upgrade to set policies by key"; + const PROMPTS_PLACEHOLDER = "Premium feature - Upgrade to set prompts by key"; + + const openModal = () => { + renderWithProviders(); + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + }; + + beforeEach(() => { + vi.mocked(getPoliciesList).mockResolvedValue({ policies: [{ policy_name: "policy-a" }] }); + vi.mocked(getPromptsList).mockResolvedValue({ prompts: [{ prompt_id: "prompt-a" }] } as any); + }); + + it("should load and offer both selectors for an admin", async () => { + openModal(); + + await waitFor(() => { + expect(screen.getByRole("option", { name: "policy-a" })).toBeInTheDocument(); + expect(screen.getByRole("option", { name: "prompt-a" })).toBeInTheDocument(); + }); + expect(getPoliciesList).toHaveBeenCalledWith("test-token"); + expect(getPromptsList).toHaveBeenCalledWith("test-token"); + expect(screen.getByPlaceholderText(POLICIES_PLACEHOLDER)).toBeInTheDocument(); + expect(screen.getByPlaceholderText(PROMPTS_PLACEHOLDER)).toBeInTheDocument(); + }); + + it("should omit both selectors and fire neither admin-only request for an internal user", async () => { + authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" }; + + openModal(); + + expect(await screen.findByTestId("org-dropdown")).toBeInTheDocument(); + + expect(getPoliciesList).not.toHaveBeenCalled(); + expect(getPromptsList).not.toHaveBeenCalled(); + expect(screen.queryByPlaceholderText(POLICIES_PLACEHOLDER)).not.toBeInTheDocument(); + expect(screen.queryByPlaceholderText(PROMPTS_PLACEHOLDER)).not.toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 189df662cd8..19abae73549 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -5,6 +5,7 @@ import { useProjects } from "@/app/(dashboard)/hooks/projects/useProjects"; import { useTags } from "@/app/(dashboard)/hooks/tags/useTags"; import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import { formatNumberWithCommas } from "@/utils/dataUtils"; import { InfoCircleOutlined } from "@ant-design/icons"; import { useQueryClient } from "@tanstack/react-query"; @@ -147,6 +148,8 @@ export const fetchUserModels = async ( const CreateKey: React.FC = ({ team, teams, data, addKey, autoOpenCreate, prefillData }) => { const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole)); + const canViewPolicies = useCan("viewPolicies"); + const canViewPrompts = useCan("viewPrompts"); const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations(); const { data: projects, isLoading: isProjectsLoading } = useProjects(); const { data: uiSettingsData } = useUISettings(); @@ -275,9 +278,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp }; fetchGuardrails(); - fetchPolicies(); - fetchPrompts(); - }, [accessToken]); + if (canViewPolicies) fetchPolicies(); + if (canViewPrompts) fetchPrompts(); + }, [accessToken, canViewPolicies, canViewPrompts]); // Fetch possible user roles when component mounts useEffect(() => { @@ -1251,74 +1254,78 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp > - - Policies{" "} - - e.stopPropagation()} // Prevent accordion from collapsing when clicking link - > - - - - - } - name="policies" - className="mt-4" - help={ - premiumUser - ? "Select existing policies or enter new ones" - : "Premium feature - Upgrade to set policies by key" - } - > - ({ value: name, label: name }))} - /> - + > + ({ value: name, label: name }))} + /> + + )} diff --git a/ui/litellm-dashboard/src/components/policies/PolicySelector.test.tsx b/ui/litellm-dashboard/src/components/policies/PolicySelector.test.tsx index 7a263743688..84dfb5cc1c5 100644 --- a/ui/litellm-dashboard/src/components/policies/PolicySelector.test.tsx +++ b/ui/litellm-dashboard/src/components/policies/PolicySelector.test.tsx @@ -7,6 +7,11 @@ import { Policy } from "./types"; vi.mock("../networking"); +const can = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useCan", () => ({ + default: (...args: unknown[]) => can(...args), +})); + const makePolicy = (overrides: Partial): Policy => ({ policy_id: "uuid-1", policy_name: "test-policy", @@ -76,6 +81,7 @@ describe("PolicySelector", () => { beforeEach(() => { vi.clearAllMocks(); + can.mockReturnValue(true); }); it("should render", () => { @@ -114,4 +120,18 @@ describe("PolicySelector", () => { renderWithProviders(); expect(networking.getPoliciesList).not.toHaveBeenCalled(); }); + + it("should render nothing and skip the admin-only fetch without the viewPolicies capability", async () => { + can.mockReturnValue(false); + vi.mocked(networking.getPoliciesList).mockResolvedValue({ policies: [] }); + + const { container } = renderWithProviders(); + + await waitFor(() => { + expect(can).toHaveBeenCalledWith("viewPolicies"); + }); + expect(networking.getPoliciesList).not.toHaveBeenCalled(); + expect(screen.queryByRole("combobox")).not.toBeInTheDocument(); + expect(container).toBeEmptyDOMElement(); + }); }); diff --git a/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx b/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx index 132538d439f..69816fafc50 100644 --- a/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx +++ b/ui/litellm-dashboard/src/components/policies/PolicySelector.tsx @@ -1,5 +1,6 @@ import React, { useEffect, useState } from "react"; import { Select } from "antd"; +import useCan from "@/app/(dashboard)/hooks/useCan"; import { Policy } from "./types"; import { getPoliciesList } from "../networking"; @@ -51,12 +52,13 @@ const PolicySelector: React.FC = ({ disabled, onPoliciesLoaded, }) => { + const canViewPolicies = useCan("viewPolicies"); const [policies, setPolicies] = useState([]); const [loading, setLoading] = useState(false); useEffect(() => { const fetchPolicies = async () => { - if (!accessToken) return; + if (!accessToken || !canViewPolicies) return; setLoading(true); try { @@ -73,12 +75,16 @@ const PolicySelector: React.FC = ({ }; fetchPolicies(); - }, [accessToken, onPoliciesLoaded]); + }, [accessToken, canViewPolicies, onPoliciesLoaded]); const handlePolicyChange = (selectedValues: string[]) => { onChange(selectedValues); }; + if (!canViewPolicies) { + return null; + } + return (
({ value: name, label: name }))} - /> - + {canViewPolicies && ( + + Policies{" "} + + e.stopPropagation()} + > + + + + + } + name="policies" + > + + + + + + + + + @@ -610,27 +602,29 @@ export function KeyEditView({ - - Policies{" "} - - - - - } - name="policies" - > - {accessToken && ( - { - form.setFieldValue("policies", v); - }} - accessToken={accessToken} - disabled={!premiumUser} - /> - )} - + {canViewPolicies && ( + + Policies{" "} + + + + + } + name="policies" + > + {accessToken && ( + { + form.setFieldValue("policies", v); + }} + accessToken={accessToken} + disabled={!premiumUser} + /> + )} + + )} 0 - ? `Current: ${keyData.metadata.prompts.join(", ")}` - : "Select or enter prompts" - } - options={promptsList.map((name) => ({ value: name, label: name }))} - /> - - + {canViewPrompts && ( + + +