diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml index 140f1155e3a..1423228e725 100644 --- a/.github/ci-coverage-allowlist.yml +++ b/.github/ci-coverage-allowlist.yml @@ -136,13 +136,6 @@ test_paths: - tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py dockerfiles: - - reason: >- - The componentized images the microservices chart deploys are built by no job; wiring both into - the scan workflow costs a full image build each and is deferred to a change that prices the - whole set - paths: - - backend/Dockerfile - - gateway/Dockerfile - reason: >- The dashboard container is a static Next.js export served by nginx, and the dashboard build and lint workflows already exercise that output, so building the image adds no signal about it diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 85291b49880..d34b0ee2e0f 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -13,6 +13,33 @@ How it solves it: - - ... +## User Flow + + + ## Relevant issues diff --git a/.github/workflows/check-schema-sync.yml b/.github/workflows/check-schema-sync.yml index 0e5e2804e60..a4e78d2c44c 100644 --- a/.github/workflows/check-schema-sync.yml +++ b/.github/workflows/check-schema-sync.yml @@ -10,6 +10,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: check-sync: name: Verify schema.prisma copies match root diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 597daebd720..02543d67a82 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -2,18 +2,19 @@ name: Check UI API Types Sync on: pull_request: - paths: - - "litellm/proxy/**" - - "litellm/types/**" - - "ui/litellm-dashboard/src/lib/http/schema.d.ts" - - "ui/litellm-dashboard/scripts/gen-api-types.mjs" - - "ui/litellm-dashboard/package.json" - - "ui/litellm-dashboard/package-lock.json" - - ".github/workflows/check-ui-api-types.yml" + branches: + - main + - litellm_internal_staging + - litellm_oss_staging + - "litellm_**" permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + jobs: check-sync: name: Verify schema.d.ts matches the proxy OpenAPI spec @@ -24,18 +25,39 @@ jobs: uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: persist-credentials: false + fetch-depth: 2 + + - name: Detect changes that can affect the generated types + id: changes + run: | + set -euo pipefail + if ! base="$(git rev-parse --verify --quiet HEAD^2 >/dev/null && git rev-parse HEAD^1)"; then + echo "Not a pull request merge commit, running the full check." + echo "relevant=true" >> "$GITHUB_OUTPUT" + exit 0 + fi + files="$(git diff --name-only "$base" HEAD)" + if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/check-ui-api-types\.yml$)' <<< "$files"; then + echo "relevant=true" >> "$GITHUB_OUTPUT" + else + echo "No proxy, types or generator changes in this pull request, nothing to verify." + echo "relevant=false" >> "$GITHUB_OUTPUT" + fi - name: Set up Python + if: steps.changes.outputs.relevant == 'true' uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 with: python-version: "3.12" - name: Set up uv + if: steps.changes.outputs.relevant == 'true' uses: ./.github/actions/setup-uv-with-retries with: version: "0.10.9" - name: Cache uv dependencies + if: steps.changes.outputs.relevant == 'true' uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 with: path: | @@ -46,14 +68,17 @@ jobs: ${{ runner.os }}-uv- - name: Install backend dependencies + 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: 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 + if: steps.changes.outputs.relevant == 'true' uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0 with: node-version-file: ui/litellm-dashboard/.nvmrc @@ -61,16 +86,19 @@ jobs: cache-dependency-path: ui/litellm-dashboard/package-lock.json - name: Install dashboard dependencies + if: steps.changes.outputs.relevant == 'true' working-directory: ui/litellm-dashboard run: npm ci - name: Regenerate types from the live spec + if: steps.changes.outputs.relevant == 'true' working-directory: ui/litellm-dashboard env: LITELLM_PYTHON: "uv run --no-sync python" run: npm run gen:api - name: Fail if types are stale + if: steps.changes.outputs.relevant == 'true' run: | if ! git diff --exit-code -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then echo "::error file=ui/litellm-dashboard/src/lib/http/schema.d.ts::Generated API types are out of sync with the proxy OpenAPI spec." diff --git a/.github/workflows/conventional-commits.yml b/.github/workflows/conventional-commits.yml index 69ade24d028..eb9eb69f8b6 100644 --- a/.github/workflows/conventional-commits.yml +++ b/.github/workflows/conventional-commits.yml @@ -14,6 +14,10 @@ on: permissions: pull-requests: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: lint-pr-title: name: Validate PR title diff --git a/.github/workflows/guard-fork-dependencies.yml b/.github/workflows/guard-fork-dependencies.yml index f4cbdd63cdf..6b366da78d4 100644 --- a/.github/workflows/guard-fork-dependencies.yml +++ b/.github/workflows/guard-fork-dependencies.yml @@ -15,6 +15,10 @@ on: permissions: {} +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: guard: name: Block fork dependency changes diff --git a/.github/workflows/helm_unit_test.yml b/.github/workflows/helm_unit_test.yml index a44d412c781..f95848945a0 100644 --- a/.github/workflows/helm_unit_test.yml +++ b/.github/workflows/helm_unit_test.yml @@ -9,6 +9,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: unit-test: runs-on: ubuntu-latest diff --git a/.github/workflows/image-scan.yml b/.github/workflows/image-scan.yml index e0c0bfcedae..8faf3ef6229 100644 --- a/.github/workflows/image-scan.yml +++ b/.github/workflows/image-scan.yml @@ -12,6 +12,11 @@ on: - docker/Dockerfile.non_root - migrations/Dockerfile - migrations/run.py + - gateway/Dockerfile + - gateway/main.py + - backend/Dockerfile + - backend/main.py + - docker/component_entrypoint.sh - litellm-proxy-extras/** - tests/proxy_migration_tests/** - uv.lock @@ -147,3 +152,63 @@ jobs: run: | python -m pip install "pytest==9.0.3" python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v + + gateway-image: + name: gateway-image + runs-on: ubuntu-latest + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Build gateway image + run: docker build -f gateway/Dockerfile -t litellm-gateway-scan:${{ github.sha }} . + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify the gateway serves offline as a non-root uid + env: + LITELLM_IMAGE: litellm-gateway-scan:${{ github.sha }} + LITELLM_COMPONENT_PORT: "4000" + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v + + backend-image: + name: backend-image + runs-on: ubuntu-latest + if: >- + github.event_name != 'pull_request' || + github.event.pull_request.head.repo.full_name == github.repository + timeout-minutes: 30 + permissions: + contents: read + steps: + - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Build backend image + run: docker build -f backend/Dockerfile -t litellm-backend-scan:${{ github.sha }} . + + - name: Set up Python + uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: "3.12" + + - name: Verify the backend serves offline as a non-root uid + env: + LITELLM_IMAGE: litellm-backend-scan:${{ github.sha }} + LITELLM_COMPONENT_PORT: "4001" + run: | + python -m pip install "pytest==9.0.3" + python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 5e333f2a3ca..3db3fb07a94 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: lint: runs-on: ubuntu-latest diff --git a/.github/workflows/test-litellm-ui-build.yml b/.github/workflows/test-litellm-ui-build.yml index 39f4bc1428a..618b0195b5a 100644 --- a/.github/workflows/test-litellm-ui-build.yml +++ b/.github/workflows/test-litellm-ui-build.yml @@ -10,6 +10,10 @@ on: - litellm_oss_staging - "litellm_**" +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: build-ui: runs-on: ubuntu-latest diff --git a/.github/workflows/test-litellm-ui-lint.yml b/.github/workflows/test-litellm-ui-lint.yml index ecc739a87e2..e03d89ee26a 100644 --- a/.github/workflows/test-litellm-ui-lint.yml +++ b/.github/workflows/test-litellm-ui-lint.yml @@ -10,6 +10,10 @@ on: - litellm_oss_staging - "litellm_**" +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: frontend-lint: runs-on: ubuntu-latest diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index a5a4e722133..05cc13d0af2 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: test: runs-on: ubuntu-latest diff --git a/.github/workflows/test-model-map.yaml b/.github/workflows/test-model-map.yaml index cf4b0eb21a1..c2770e5da4c 100644 --- a/.github/workflows/test-model-map.yaml +++ b/.github/workflows/test-model-map.yaml @@ -11,6 +11,10 @@ on: permissions: contents: read +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + jobs: validate-model-prices-json: runs-on: ubuntu-latest diff --git a/.github/workflows/test-unit-proxy-endpoints.yml b/.github/workflows/test-unit-proxy-endpoints.yml index 20b89c72440..2ea3c521e8b 100644 --- a/.github/workflows/test-unit-proxy-endpoints.yml +++ b/.github/workflows/test-unit-proxy-endpoints.yml @@ -38,12 +38,15 @@ jobs: tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/batches_endpoints + tests/test_litellm/proxy/fine_tuning_endpoints + tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints tests/test_litellm/proxy/response_api_endpoints tests/test_litellm/proxy/image_endpoints tests/test_litellm/proxy/vector_store_endpoints tests/test_litellm/proxy/agent_endpoints tests/test_litellm/proxy/a2a + tests/test_litellm/proxy/credential_endpoints tests/test_litellm/proxy/discovery_endpoints tests/test_litellm/proxy/health_endpoints tests/test_litellm/proxy/shutdown diff --git a/CLAUDE.md b/CLAUDE.md index 59929143c46..b0da2970fc8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -41,7 +41,7 @@ Python max line length is 120, not 88 When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing -`make pre-commit` saves its complete output to a log file in .git (overwriting previous pre-commit logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice +`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in diff --git a/Makefile b/Makefile index 493828571b7..94d8c875af5 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \ lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \ install-dev install-proxy-dev install-test-deps install-hooks \ - install-helm-unittest check-circular-imports check-import-safety pre-commit \ + install-helm-unittest check-circular-imports check-import-safety check pre-commit \ lint-install lint-fetch-base bootstrap # Default target @@ -22,7 +22,8 @@ help: @echo " make install-test-deps - Install the full local test environment" @echo " make install-helm-unittest - Install helm unittest plugin" @echo " make install-hooks - Install git hooks (Conventional Commits + Branches)" - @echo " make pre-commit - Run CI-equivalent lint on staged files (run before committing)" + @echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged" + @echo " make pre-commit - Legacy alias for make check" @echo " make format - Apply ruff format code formatting" @echo " make format-check - Check ruff format code formatting (matches CI)" @echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)" @@ -236,13 +237,20 @@ lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline # Faster linting for local development (only checks changed code) lint-dev: lint-format-changed check-circular-imports check-import-safety -# Run the gating CI checks against your staged files right before committing. Mirrors +# Run the gating CI checks against your changes. Scopes to staged files when anything +# is staged (warning about changed files left unstaged); with nothing staged it falls +# back to the working tree's diff against the merge base with the base branch, so a +# fresh merge commit or an unstaged working tree still gets checked. Mirrors # test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and -# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage. +# check-ui-api-types.yml (API-type drift), skipping any whose files aren't in scope. # Not auto-installed as a git hook so it never slows an unrelated human commit. -pre-commit: bootstrap +check: bootstrap ./scripts/pre_commit_lint.sh +pre-commit: + @echo "make pre-commit is a legacy alias; use make check" >&2 + @$(MAKE) check + # Testing targets test: install-test-deps $(UV_RUN) pytest tests/ diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 632743236c4..0385f7a96e7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 28842 + "limit": 27731 }, "reportArgumentType": { - "limit": 2634 + "limit": 2626 }, "reportAssignmentType": { "limit": 329 @@ -12,7 +12,7 @@ "limit": 514 }, "reportCallIssue": { - "limit": 117 + "limit": 116 }, "reportConstantRedefinition": { "limit": 40 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 9103 + "limit": 8807 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5843 + "limit": 5835 }, "reportMissingTypeArgument": { - "limit": 15816 + "limit": 15790 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1078 + "limit": 1077 }, "reportOptionalOperand": { "limit": 0 @@ -84,34 +84,34 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1825 + "limit": 1824 }, "reportRedeclaration": { "limit": 8 }, "reportReturnType": { - "limit": 218 + "limit": 217 }, "reportTypedDictNotRequiredAccess": { - "limit": 27 + "limit": 26 }, "reportUndefinedVariable": { "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45110 + "limit": 45063 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39838 + "limit": 39773 }, "reportUnknownParameterType": { - "limit": 20237 + "limit": 20207 }, "reportUnknownVariableType": { - "limit": 31383 + "limit": 31281 }, "reportUnnecessaryCast": { "limit": 122 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 864 + "limit": 862 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 0036603bcd1..f0914240f79 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -3,11 +3,25 @@ import base64 import json +from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Final, + List, + Literal, + Optional, + Protocol, + TypedDict, + Union, + cast, +) from uuid import NAMESPACE_URL, uuid5 from fastapi import HTTPException +from pydantic import ValidationError import litellm from litellm import Router, verbose_logger @@ -31,9 +45,12 @@ from litellm.proxy._types import ( ) from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + apply_unified_file_ids, + ensure_batch_response_managed_file_ids, get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, + map_raw_file_ids_to_unified, normalize_mime_type_for_provider, resolve_managed_output_file_model_name, ) @@ -61,6 +78,9 @@ if TYPE_CHECKING: if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from prisma.models import ( + LiteLLM_ManagedObjectTable as PrismaManagedObjectRow, + ) from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache from litellm.proxy.utils import PrismaClient as _PrismaClient @@ -74,11 +94,92 @@ else: PrismaClient = Any +def _sanitized_parse_error(e: Exception) -> str: + return ( + str(e.errors(include_input=False, include_url=False, include_context=False)) + if isinstance(e, ValidationError) + else type(e).__name__ + ) + + +def _decode_json_blob(blob: object) -> object: + return json.loads(blob) if isinstance(blob, str) else blob + + +def _parse_managed_batch_row(row: "PrismaManagedObjectRow") -> Optional[LiteLLMBatch]: + try: + batch_obj: Final = LiteLLMBatch.model_validate(_decode_json_blob(row.file_object)) + except Exception as e: + verbose_logger.warning(f"Failed to parse batch object {row.unified_object_id}: {_sanitized_parse_error(e)}") + return None + batch_obj.id = row.unified_object_id + return batch_obj + + +def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) -> Optional[OpenAIFileObject]: + if raw_file_object is None: + return None + try: + return OpenAIFileObject.model_validate(raw_file_object) + except Exception as e: + verbose_logger.warning(f"Failed to parse managed file object {unified_file_id}: {_sanitized_parse_error(e)}") + return None + + +class _ManagedFileRow(Protocol): + unified_file_id: str + file_object: OpenAIFileObject + storage_backend: Optional[str] + storage_url: Optional[str] + created_by: Optional[str] + team_id: Optional[str] + + def model_dump(self) -> Mapping[str, object]: ... + + +class _ManagedFileTableActions(Protocol): + async def find_first(self, where: Mapping[str, object]) -> Optional[_ManagedFileRow]: ... + + async def find_many(self, where: Mapping[str, object]) -> Sequence[_ManagedFileRow]: ... + + async def upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]]) -> _ManagedFileRow: ... + + async def delete(self, where: Mapping[str, str]) -> Optional[_ManagedFileRow]: ... + + +class _ManagedObjectTableActions(Protocol): + async def find_first(self, where: Mapping[str, object]) -> "Optional[PrismaManagedObjectRow]": ... + + async def find_many( + self, + where: Mapping[str, object], + take: int, + order: Union[Mapping[str, str], Sequence[Mapping[str, str]]], + cursor: Mapping[str, str] = ..., + skip: int = ..., + ) -> "Sequence[PrismaManagedObjectRow]": ... + + async def upsert( + self, where: Mapping[str, str], data: Mapping[str, Mapping[str, object]] + ) -> "PrismaManagedObjectRow": ... + + +class _CursorPageArgs(TypedDict, total=False): + cursor: Mapping[str, str] + skip: int + + +def _managed_file_table(prisma_client: PrismaClient) -> _ManagedFileTableActions: + return prisma_client.db.litellm_managedfiletable + + +def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableActions: + return prisma_client.db.litellm_managedobjecttable + + class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Class variables or attributes - def __init__( - self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient - ): + def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client @@ -97,9 +198,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings: Dict[str, str], user_api_key_dict: UserAPIKeyAuth, ) -> None: - verbose_logger.info( - f"Storing LiteLLM Managed File object with id={file_id} in cache" - ) + verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache") if file_object is not None: litellm_managed_file_object = LiteLLM_ManagedFileTable( unified_file_id=file_id, @@ -150,13 +249,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): f"storage_url={db_data.get('storage_url')}" ) - result = await self.prisma_client.db.litellm_managedfiletable.upsert( + result = await _managed_file_table(self.prisma_client).upsert( where={"unified_file_id": file_id}, data={"create": db_data, "update": update_data}, ) - verbose_logger.debug( - f"LiteLLM Managed File object with id={file_id} stored in db: {result}" - ) + verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") async def store_unified_object_id( self, @@ -167,9 +264,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, ) -> None: - verbose_logger.info( - f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache" - ) + 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, model_object_id=model_object_id, @@ -182,7 +277,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): litellm_parent_otel_span=litellm_parent_otel_span, ) - await self.prisma_client.db.litellm_managedobjecttable.upsert( + await _managed_object_table(self.prisma_client).upsert( where={"unified_object_id": unified_object_id}, data={ "create": { @@ -219,9 +314,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return LiteLLM_ManagedFileTable.model_validate(result) ## CHECK DB - db_object = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + db_object = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if db_object: return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump()) @@ -231,9 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self, file_id: str, litellm_parent_otel_span: Optional[Span] = None ) -> OpenAIFileObject: ## get old value - initial_value = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + initial_value = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if initial_value is None: raise Exception(f"LiteLLM Managed File object with id={file_id} not found") ## delete old value @@ -242,15 +333,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): value=None, litellm_parent_otel_span=litellm_parent_otel_span, ) - await self.prisma_client.db.litellm_managedfiletable.delete( - where={"unified_file_id": file_id} - ) + await _managed_file_table(self.prisma_client).delete(where={"unified_file_id": file_id}) return initial_value.file_object - async def can_user_call_unified_file_id( - self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: - managed_file = await self.prisma_client.db.litellm_managedfiletable.find_first( + async def can_user_call_unified_file_id(self, unified_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + managed_file = await _managed_file_table(self.prisma_client).find_first( where={"unified_file_id": unified_file_id} ) @@ -265,13 +352,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"File not found: {unified_file_id}", ) - async def can_user_call_unified_object_id( - self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth - ) -> bool: - managed_object = ( - await self.prisma_client.db.litellm_managedobjecttable.find_first( - where={"unified_object_id": unified_object_id} - ) + async def can_user_call_unified_object_id(self, unified_object_id: str, user_api_key_dict: UserAPIKeyAuth) -> bool: + managed_object = await _managed_object_table(self.prisma_client).find_first( + where={"unified_object_id": unified_object_id} ) if managed_object: @@ -293,34 +376,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): provider: Optional[str] = None, target_model_names: Optional[str] = None, llm_router: Optional[Router] = None, - ) -> Dict[str, Any]: + ) -> Dict[str, object]: # Provider filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the provider information # To support provider filtering, we would need to store the provider information in the encoded object ids if provider: - raise Exception( - "Filtering by 'provider' is not supported when using managed batches." - ) + raise Exception("Filtering by 'provider' is not supported when using managed batches.") # Model name filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the model name # A hash of the model name + litellm_params for the model name is encoded as the model id. This is not sufficient to reliably map the target model names to the model ids. if target_model_names: - raise Exception( - "Filtering by 'target_model_names' is not supported when using managed batches." - ) + raise Exception("Filtering by 'target_model_names' is not supported when using managed batches.") owner_filter = build_owner_filter(user_api_key_dict) if owner_filter is None: return build_list_page([]) - where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter} + where_clause: Dict[str, object] = {"file_purpose": "batch", **owner_filter} if after: - cursor_row = ( - await self.prisma_client.db.litellm_managedobjecttable.find_first( - where={**where_clause, "unified_object_id": after} - ) + cursor_row = await _managed_object_table(self.prisma_client).find_first( + where={**where_clause, "unified_object_id": after} ) if cursor_row is None: raise HTTPException( @@ -328,12 +405,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): detail=f"Invalid 'after' cursor: no batch found with id '{after}'.", ) - page_size = limit or 20 - cursor_args: Dict[str, Any] = ( - {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - ) + page_size: Final = min(limit or 20, 100) + cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( + batches = await _managed_object_table(self.prisma_client).find_many( where=where_clause, take=page_size + 1, order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], @@ -342,25 +417,54 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): has_more = len(batches) > page_size - batch_objects: List[LiteLLMBatch] = [] - for batch in batches[:page_size]: - try: - batch_data = ( - json.loads(batch.file_object) - if isinstance(batch.file_object, str) - else batch.file_object - ) - batch_obj = LiteLLMBatch.model_validate(batch_data) - batch_obj.id = batch.unified_object_id - batch_objects.append(batch_obj) + parsed_rows: Final = tuple( + (row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None + ) + unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified( + raw_file_ids=frozenset( + file_id + for _, batch_obj in parsed_rows + for file_id in (batch_obj.input_file_id, batch_obj.output_file_id, batch_obj.error_file_id) + if file_id and not _is_base64_encoded_unified_file_id(file_id) + ), + prisma_client=self.prisma_client, + ) + resolved_batches: Final = [ + await self._resolve_listed_batch( + row=row, + batch_obj=batch_obj, + unified_id_by_raw_id=unified_id_by_raw_id, + user_api_key_dict=user_api_key_dict, + ) + for row, batch_obj in parsed_rows + ] + return build_list_page( + [batch_obj for batch_obj in resolved_batches if batch_obj is not None], + has_more=has_more, + ) - except Exception as e: - verbose_logger.warning( - f"Failed to parse batch object {batch.unified_object_id}: {e}" - ) - continue - - return build_list_page(batch_objects, has_more=has_more) + async def _resolve_listed_batch( + self, + row: "PrismaManagedObjectRow", + batch_obj: LiteLLMBatch, + unified_id_by_raw_id: Mapping[str, str], + user_api_key_dict: UserAPIKeyAuth, + ) -> Optional[LiteLLMBatch]: + apply_unified_file_ids(batch_obj, unified_id_by_raw_id) + try: + await ensure_batch_response_managed_file_ids( + response=batch_obj, + managed_files_obj=self, + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_logger, + user_api_key_dict=user_api_key_dict, + db_batch_object=row, + unified_batch_id=_is_base64_encoded_unified_file_id(row.unified_object_id), + ) + except Exception as e: + verbose_logger.warning(f"Failed to resolve managed file ids for batch {row.unified_object_id}: {e}") + return None + return batch_obj async def get_user_created_file_ids( self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str] @@ -377,33 +481,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if owner_filter is None: return [] - file_ids = await self.prisma_client.db.litellm_managedfiletable.find_many( + file_ids = await _managed_file_table(self.prisma_client).find_many( where={ **owner_filter, "flat_model_file_ids": {"hasSome": model_object_ids}, } ) return [ - OpenAIFileObject.model_validate(row.file_object).model_copy( - update={"id": row.unified_file_id} - ) + parsed_file_object.model_copy(update={"id": row.unified_file_id}) for row in file_ids - if row.file_object is not None + if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None ] - async def check_managed_file_id_access( - self, data: Dict, user_api_key_dict: UserAPIKeyAuth - ) -> bool: + async def check_managed_file_id_access(self, data: Dict, user_api_key_dict: UserAPIKeyAuth) -> bool: retrieve_file_id = cast(Optional[str], data.get("file_id")) - potential_file_id = ( - _is_base64_encoded_unified_file_id(retrieve_file_id) - if retrieve_file_id - else False - ) + potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and retrieve_file_id: - if await self.can_user_call_unified_file_id( - retrieve_file_id, user_api_key_dict - ): + if await self.can_user_call_unified_file_id(retrieve_file_id, user_api_key_dict): return True else: raise HTTPException( @@ -412,9 +506,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) return False - async def check_file_ids_access( - self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth - ) -> None: + async def check_file_ids_access(self, file_ids: List[str], user_api_key_dict: UserAPIKeyAuth) -> None: """ Check if the user has access to a list of file IDs. Only checks managed (unified) file IDs. @@ -429,9 +521,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for file_id in file_ids: is_unified_file_id = _is_base64_encoded_unified_file_id(file_id) if is_unified_file_id: - if not await self.can_user_call_unified_file_id( - file_id, user_api_key_dict - ): + if not await self.can_user_call_unified_file_id(file_id, user_api_key_dict): raise HTTPException( status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}", @@ -459,10 +549,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ### HANDLE TRANSFORMATIONS ### # Check both completion and acompletion call types - is_completion_call = ( - call_type == CallTypes.completion.value - or call_type == CallTypes.acompletion.value - ) + is_completion_call = call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value if is_completion_call: messages = data.get("messages") @@ -475,9 +562,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check if any files are stored in storage backends and need base64 conversion # This is needed for Vertex AI/Gemini which requires base64 content - is_vertex_ai = model and ( - "vertex_ai" in model or "gemini" in model.lower() - ) + is_vertex_ai = model and ("vertex_ai" in model or "gemini" in model.lower()) if is_vertex_ai: await self._convert_storage_files_to_base64( messages=messages, @@ -489,10 +574,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids, user_api_key_dict.parent_otel_span ) data["model_file_id_mapping"] = model_file_id_mapping - elif ( - call_type == CallTypes.aresponses.value - or call_type == CallTypes.responses.value - ): + elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value: # Handle managed files in responses API input and tools file_ids = [] @@ -519,23 +601,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if tools: unified_vs_ids = self.get_vector_store_ids_from_file_search_tools(tools) if unified_vs_ids: - await self.check_vector_store_ids_access( - unified_vs_ids, user_api_key_dict - ) + await self.check_vector_store_ids_access(unified_vs_ids, user_api_key_dict) elif call_type == CallTypes.afile_content.value: retrieve_file_id = cast(Optional[str], data.get("file_id")) - potential_file_id = ( - _is_base64_encoded_unified_file_id(retrieve_file_id) - if retrieve_file_id - else False - ) + potential_file_id = _is_base64_encoded_unified_file_id(retrieve_file_id) if retrieve_file_id else False if potential_file_id and "llm_output_file_id," in potential_file_id: model_id = self.get_model_id_from_unified_file_id(potential_file_id) if model_id: data["model"] = model_id - data["file_id"] = self.get_output_file_id_from_unified_file_id( - potential_file_id - ) + data["file_id"] = self.get_output_file_id_from_unified_file_id(potential_file_id) elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: @@ -552,10 +626,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ): accessor_key: Optional[str] = None retrieve_object_id: Optional[str] = None - if ( - call_type == CallTypes.aretrieve_batch.value - or call_type == CallTypes.acancel_batch.value - ): + if call_type == CallTypes.aretrieve_batch.value or call_type == CallTypes.acancel_batch.value: accessor_key = "batch_id" elif ( call_type == CallTypes.acancel_fine_tuning_job.value @@ -567,32 +638,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): retrieve_object_id = cast(Optional[str], data.get(accessor_key)) potential_llm_object_id = ( - _is_base64_encoded_unified_file_id(retrieve_object_id) - if retrieve_object_id - else False + _is_base64_encoded_unified_file_id(retrieve_object_id) if retrieve_object_id else False ) if potential_llm_object_id and retrieve_object_id: ## VALIDATE USER HAS ACCESS TO THE OBJECT ## - if not await self.can_user_call_unified_object_id( - retrieve_object_id, user_api_key_dict - ): + if not await self.can_user_call_unified_object_id(retrieve_object_id, user_api_key_dict): raise HTTPException( status_code=403, detail=f"User {user_api_key_dict.user_id} does not have access to the object {retrieve_object_id}", ) ## for managed batch id - get the model id - potential_model_id = get_model_id_from_unified_batch_id( - potential_llm_object_id - ) + potential_model_id = get_model_id_from_unified_batch_id(potential_llm_object_id) if potential_model_id is None: raise Exception( f"LiteLLM Managed {accessor_key} with id={retrieve_object_id} is invalid - does not contain encoded model_id." ) data["model"] = potential_model_id - data[accessor_key] = get_batch_id_from_unified_batch_id( - potential_llm_object_id - ) + data[accessor_key] = get_batch_id_from_unified_batch_id(potential_llm_object_id) elif call_type == CallTypes.acreate_fine_tuning_job.value: input_file_id = cast(Optional[str], data.get("training_file")) if input_file_id: @@ -648,24 +711,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if accessor_key: input_file_id = cast(Optional[str], kwargs.get(accessor_key)) - model_file_id_mapping = cast( - Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping") - ) + model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")) # model_info may be at top-level or nested under litellm_metadata # (batch/file operations use litellm_metadata) model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None)) if model_id is None: model_id = cast( Optional[str], - kwargs.get("litellm_metadata", {}) - .get("model_info", {}) - .get("id", None), + kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None), ) mapped_file_id: Optional[str] = None if input_file_id and model_file_id_mapping and model_id: - mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get( - model_id, None - ) + mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None) if mapped_file_id: kwargs[accessor_key] = mapped_file_id @@ -691,9 +748,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_input( - self, input: Union[str, List[Dict[str, Any]]] - ) -> List[str]: + def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]: """ Gets file ids from responses API input. @@ -725,19 +780,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): content = item.get("content") if isinstance(content, list): for content_item in content: - if ( - isinstance(content_item, dict) - and content_item.get("type") == "input_file" - ): + if isinstance(content_item, dict) and content_item.get("type") == "input_file": file_id = content_item.get("file_id") if file_id: file_ids.append(file_id) return file_ids - def get_file_ids_from_responses_tools( - self, tools: List[Dict[str, Any]] - ) -> List[str]: + def get_file_ids_from_responses_tools(self, tools: List[Dict[str, object]]) -> List[str]: """ Gets file ids from responses API tools parameter. @@ -770,9 +820,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return file_ids - def get_vector_store_ids_from_file_search_tools( - self, tools: List[Dict[str, Any]] - ) -> List[str]: + def get_vector_store_ids_from_file_search_tools(self, tools: List[Dict[str, object]]) -> List[str]: """ Extract unified vector_store_ids from file_search tools. @@ -865,9 +913,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ), ) - async def get_model_file_id_mapping( - self, file_ids: List[str], litellm_parent_otel_span: Span - ) -> dict: + async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict: """ Get model-specific file IDs for a list of proxy file IDs. Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id @@ -897,9 +943,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Get all cache keys matching the pattern file_id:* for file_id in litellm_managed_file_ids: # Search for any cache key starting with this file_id - unified_file_object = await self.get_unified_file_id( - file_id, litellm_parent_otel_span - ) + unified_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) if unified_file_object: file_id_mapping[file_id] = unified_file_object.model_mappings @@ -917,9 +961,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): raise Exception("LLM Router not initialized. Ensure models added to proxy.") responses = [] for model in target_model_names_list: - individual_response = await llm_router.acreate_file( - model=model, **_create_file_request - ) + individual_response = await llm_router.acreate_file(model=model, **_create_file_request) responses.append(individual_response) return responses @@ -950,9 +992,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings: Dict[str, str] = {} for file_object in responses: - model_file_id_mapping = file_object._hidden_params.get( - "model_file_id_mapping" - ) + model_file_id_mapping = file_object._hidden_params.get("model_file_id_mapping") if model_file_id_mapping and isinstance(model_file_id_mapping, dict): model_mappings.update(model_file_id_mapping) @@ -967,17 +1007,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Emit Prometheus metrics for managed file creation prom_logger = self._get_prometheus_logger() if prom_logger: - first_model = ( - target_model_names_list[0] if target_model_names_list else None - ) + first_model = target_model_names_list[0] if target_model_names_list else None first_provider = "" if responses: - first_provider = ( - getattr(responses[0], "_hidden_params", {}).get( - "custom_llm_provider" - ) - or "" - ) + first_provider = getattr(responses[0], "_hidden_params", {}).get("custom_llm_provider") or "" prom_logger.record_managed_file_created( model=first_model or "", api_provider=first_provider, @@ -1020,9 +1053,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) # Convert to URL-safe base64 and strip padding - base64_unified_file_id = ( - base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") - ) + base64_unified_file_id = base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") ## CREATE RESPONSE OBJECT @@ -1039,46 +1070,26 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response - def get_unified_generic_response_id( - self, model_id: str, generic_response_id: str - ) -> str: - unified_generic_response_id = ( - SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format( - model_id, generic_response_id - ) - ) - return ( - base64.urlsafe_b64encode(unified_generic_response_id.encode()) - .decode() - .rstrip("=") + def get_unified_generic_response_id(self, model_id: str, generic_response_id: str) -> str: + unified_generic_response_id = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format( + model_id, generic_response_id ) + return base64.urlsafe_b64encode(unified_generic_response_id.encode()).decode().rstrip("=") def get_unified_batch_id(self, batch_id: str, model_id: str) -> str: - unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( - model_id, batch_id - ) + unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") - def get_unified_output_file_id( - self, output_file_id: str, model_id: str, model_name: Optional[str] - ) -> str: - deterministic_uuid: Final = uuid5( - uuid5(NAMESPACE_URL, model_id), output_file_id - ) - unified_output_file_id = ( - SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - "application/json", - str(deterministic_uuid), - model_name or "", - output_file_id, - model_id, - ) - ) - return ( - base64.urlsafe_b64encode(unified_output_file_id.encode()) - .decode() - .rstrip("=") + def get_unified_output_file_id(self, output_file_id: str, model_id: str, model_name: Optional[str]) -> str: + deterministic_uuid: Final = uuid5(uuid5(NAMESPACE_URL, model_id), output_file_id) + unified_output_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", + str(deterministic_uuid), + model_name or "", + output_file_id, + model_id, ) + return base64.urlsafe_b64encode(unified_output_file_id.encode()).decode().rstrip("=") def get_model_id_from_unified_file_id(self, file_id: str) -> str: return file_id.split("llm_output_file_model_id,")[1].split(";")[0] @@ -1086,59 +1097,39 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): def get_output_file_id_from_unified_file_id(self, file_id: str) -> str: marker = "llm_output_file_id," if marker not in file_id: - raise ValueError( - f"Unified id does not contain {marker!r}: {file_id[:80]!r}" - ) + raise ValueError(f"Unified id does not contain {marker!r}: {file_id[:80]!r}") return file_id.split(marker, 1)[1].split(";")[0] async def async_post_call_success_hook( self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes - ) -> Any: + ) -> LLMResponseTypes: if isinstance(response, LiteLLMBatch): ## Check if unified_file_id is in the response - unified_file_id = response._hidden_params.get( - "unified_file_id" - ) # managed file id - unified_batch_id = response._hidden_params.get( - "unified_batch_id" - ) # managed batch id + unified_file_id = response._hidden_params.get("unified_file_id") # managed file id + unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id model_id = cast(Optional[str], response._hidden_params.get("model_id")) model_name = cast(Optional[str], response._hidden_params.get("model_name")) resolved_model_name = resolve_managed_output_file_model_name( - unified_input_file_id=unified_file_id - if isinstance(unified_file_id, str) - else response.input_file_id, + unified_input_file_id=unified_file_id if isinstance(unified_file_id, str) else response.input_file_id, fallback_model_name=model_name, ) original_response_id = response.id if (unified_batch_id or unified_file_id) and model_id: - response.id = self.get_unified_batch_id( - batch_id=response.id, model_id=model_id - ) + response.id = self.get_unified_batch_id(batch_id=response.id, model_id=model_id) # Handle both output_file_id and error_file_id for file_attr in ["output_file_id", "error_file_id"]: file_id_value = getattr(response, file_attr, None) if file_id_value and model_id: - decoded_output_file_id = _is_base64_encoded_unified_file_id( - file_id_value - ) - if ( - decoded_output_file_id - and "llm_output_file_id," in decoded_output_file_id - ): - provider_file_id = ( - self.get_output_file_id_from_unified_file_id( - decoded_output_file_id - ) - ) + decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value) + if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id: + provider_file_id = self.get_output_file_id_from_unified_file_id(decoded_output_file_id) unified_file_id = file_id_value elif decoded_output_file_id: verbose_logger.warning( - f"Skipping {file_attr}={file_id_value!r}: " - "unified id is not a managed file output id" + f"Skipping {file_attr}={file_id_value!r}: unified id is not a managed file output id" ) continue else: @@ -1157,23 +1148,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Import module and use getattr for better testability with mocks import litellm.proxy.proxy_server as proxy_server_module - _llm_router = getattr( - proxy_server_module, "llm_router", None - ) + _llm_router = getattr(proxy_server_module, "llm_router", None) if _llm_router is not None and model_id: - _creds = ( - _llm_router.get_deployment_credentials_with_provider( - model_id - ) - or {} - ) + _creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {} file_object = await litellm.afile_retrieve( file_id=provider_file_id, **_creds, ) else: file_object = await litellm.afile_retrieve( - custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", # type: ignore[arg-type] + custom_llm_provider=model_name.split("/")[0] + if model_name and "/" in model_name + else "openai", # type: ignore[arg-type] file_id=provider_file_id, ) verbose_logger.debug( @@ -1227,9 +1213,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): elif isinstance(response, LiteLLMFineTuningJob): ## Check if unified_file_id is in the response - unified_file_id = response._hidden_params.get( - "unified_file_id" - ) # managed file id + unified_file_id = response._hidden_params.get("unified_file_id") # managed file id unified_finetuning_job_id = response._hidden_params.get( "unified_finetuning_job_id" ) # managed finetuning job id @@ -1237,9 +1221,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_name = cast(Optional[str], response._hidden_params.get("model_name")) original_response_id = response.id if (unified_file_id or unified_finetuning_job_id) and model_id: - response.id = self.get_unified_generic_response_id( - model_id=model_id, generic_response_id=response.id - ) + response.id = self.get_unified_generic_response_id(model_id=model_id, generic_response_id=response.id) await self.store_unified_object_id( unified_object_id=response.id, file_object=response, @@ -1254,9 +1236,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ ## check if file object if hasattr(response, "data") and isinstance(response.data, list): - if all( - isinstance(file_object, FileObject) for file_object in response.data - ): + if all(isinstance(file_object, FileObject) for file_object in response.data): ## Get all file id's ## Check which file id's were created by the user ## Filter the response to only include the files created by the user @@ -1265,9 +1245,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_object.id for file_object in cast(List[FileObject], response.data) # type: ignore ] - user_created_file_ids = await self.get_user_created_file_ids( - user_api_key_dict, file_ids - ) + 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 return response @@ -1275,11 +1253,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response async def afile_retrieve( - self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router=None + self, file_id: str, litellm_parent_otel_span: Optional[Span], llm_router: Optional[Router] = None ) -> OpenAIFileObject: - stored_file_object = await self.get_unified_file_id( - file_id, litellm_parent_otel_span - ) + stored_file_object = await self.get_unified_file_id(file_id, litellm_parent_otel_span) # Case 1 : This is not a managed file if not stored_file_object: @@ -1302,21 +1278,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) try: - model_id, model_file_id = next( - iter(stored_file_object.model_mappings.items()) - ) - credentials = ( - llm_router.get_deployment_credentials_with_provider(model_id) or {} - ) - response = await litellm.afile_retrieve( - file_id=model_file_id, **credentials - ) + model_id, model_file_id = next(iter(stored_file_object.model_mappings.items())) + credentials = llm_router.get_deployment_credentials_with_provider(model_id) or {} + response = await litellm.afile_retrieve(file_id=model_file_id, **credentials) response.id = file_id # Replace with unified ID return response except Exception as e: - raise Exception( - f"Failed to retrieve file {file_id} from provider: {str(e)}" - ) from e + raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e async def afile_list( self, @@ -1353,12 +1321,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return False except Exception as e: - verbose_logger.warning( - f"Error checking batch polling configuration: {e}. Assuming disabled." - ) + verbose_logger.warning(f"Error checking batch polling configuration: {e}. Assuming disabled.") return False - async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, Any]]: + async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, object]]: """ Find batches that reference this file and still need cost tracking. Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost. @@ -1374,9 +1340,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Get model-specific file IDs for this unified file ID if it's a managed file try: - model_file_id_mapping = await self.get_model_file_id_mapping( - [file_id], litellm_parent_otel_span=None - ) + model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span=None) if model_file_id_mapping and file_id in model_file_id_mapping: # Add all provider file IDs for this unified file @@ -1384,8 +1348,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_ids_to_check.extend(provider_file_ids) except Exception as e: verbose_logger.debug( - f"Could not get model file ID mapping for {file_id}: {e}. " - f"Will only check unified file ID." + f"Could not get model file ID mapping for {file_id}: {e}. Will only check unified file ID." ) MAX_MATCHES_TO_RETURN = 10 @@ -1403,11 +1366,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): for batch in batches: try: # Parse the batch file_object to check for file references - batch_data = ( - json.loads(batch.file_object) - if isinstance(batch.file_object, str) - else batch.file_object - ) + batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object # Extract file IDs from batch # Batches typically reference the unified file ID in input_file_id @@ -1416,9 +1375,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): output_file_id = batch_data.get("output_file_id") error_file_id = batch_data.get("error_file_id") - referenced_file_ids = [ - fid for fid in [input_file_id, output_file_id, error_file_id] if fid - ] + referenced_file_ids = [fid for fid in [input_file_id, output_file_id, error_file_id] if fid] # Check if any referenced file ID matches the file we're trying to delete if any(ref_id in file_ids_to_check for ref_id in referenced_file_ids): @@ -1430,9 +1387,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): } ) except Exception as e: - verbose_logger.warning( - f"Error parsing batch object {batch.unified_object_id}: {e}" - ) + verbose_logger.warning(f"Error parsing batch object {batch.unified_object_id}: {e}") continue return referencing_batches @@ -1461,21 +1416,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if referencing_batches: # File is referenced by non-terminal batches and polling is enabled - MAX_BATCHES_IN_ERROR = ( - 5 # Limit batches shown in error message for readability - ) + MAX_BATCHES_IN_ERROR = 5 # Limit batches shown in error message for readability # Show up to MAX_BATCHES_IN_ERROR in the error message batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR] - batch_statuses = [ - f"{b['batch_id']}: {b['status']}" for b in batches_to_show - ] + batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show] # Determine the count message count_message = f"{len(referencing_batches)}" - if ( - len(referencing_batches) >= 10 - ): # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file + if len(referencing_batches) >= 10: # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file count_message = "10+" error_message = ( @@ -1516,23 +1465,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): await self._check_file_deletion_allowed(file_id) # file_id = convert_b64_uid_to_unified_uid(file_id) - model_file_id_mapping = await self.get_model_file_id_mapping( - [file_id], litellm_parent_otel_span - ) + model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) delete_response = None specific_model_file_id_mapping = model_file_id_mapping.get(file_id) if specific_model_file_id_mapping: # Remove conflicting keys from data to avoid duplicate keyword arguments - filtered_data = { - k: v for k, v in data.items() if k not in ("model", "file_id") - } + filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")} for model_id, model_file_id in specific_model_file_id_mapping.items(): delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore - stored_file_object = await self.delete_unified_file_id( - file_id, litellm_parent_otel_span - ) + stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span) # Record successful deletion metric only on actual success if stored_file_object or delete_response: @@ -1559,9 +1502,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): Get the content of a file from first model that has it """ model_file_id_mapping = data.pop("model_file_id_mapping", None) - model_file_id_mapping = ( - model_file_id_mapping - or await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span) + model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping( + [file_id], litellm_parent_otel_span ) specific_model_file_id_mapping = model_file_id_mapping.get(file_id) @@ -1574,13 +1516,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # against the deployment's configured bucket, which they only # trust from this immutable server-side snapshot, never from # request params. - credentials = llm_router.get_deployment_credentials_with_provider( - model_id=model_id - ) + credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id) if credentials is not None: - data["_litellm_internal_model_credentials"] = cast( - Dict, MappingProxyType(dict(credentials)) - ) + data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials))) else: data.pop("_litellm_internal_model_credentials", None) return await llm_router.afile_content(model=model_id, file_id=provider_file_id, **data) # type: ignore @@ -1615,9 +1553,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Check database for storage backend info # IMPORTANT: The database stores the base64 encoded unified_file_id (not the decoded version) # So we query with the original file_id (which is base64 encoded) - db_file = await self.prisma_client.db.litellm_managedfiletable.find_first( - where={"unified_file_id": file_id} - ) + db_file = await _managed_file_table(self.prisma_client).find_first(where={"unified_file_id": file_id}) if not db_file or not db_file.storage_backend or not db_file.storage_url: continue @@ -1643,22 +1579,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_content = await storage_backend.download_file(storage_url) # Determine content type from file object - content_type = self._get_content_type_from_file_object( - db_file.file_object - ) + content_type = self._get_content_type_from_file_object(db_file.file_object) # Convert to base64 base64_data = base64.b64encode(file_content).decode("utf-8") base64_data_uri = f"data:{content_type};base64,{base64_data}" # Update messages to use base64 instead of file_id - self._update_messages_with_base64_data( - messages, file_id, base64_data_uri, content_type - ) + self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type) except Exception as e: - verbose_logger.exception( - f"Error converting file {file_id} from storage backend to base64: {str(e)}" - ) + verbose_logger.exception(f"Error converting file {file_id} from storage backend to base64: {str(e)}") # Continue with other files even if one fails continue diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 1f693526d1f..66fac8d76ee 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -19,6 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import delete_cached_project_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field from litellm.proxy.management_helpers.utils import ( @@ -514,6 +515,7 @@ async def update_project( litellm_proxy_admin_name, premium_user, prisma_client, + user_api_key_cache, ) try: @@ -672,6 +674,11 @@ async def update_project( include={"litellm_budget_table": True, "object_permission": True}, ) + await delete_cached_project_object( + project_id=data.project_id, + user_api_key_cache=user_api_key_cache, + ) + return updated_project except Exception as e: verbose_proxy_logger.exception( @@ -710,7 +717,7 @@ async def delete_project( }' ``` """ - from litellm.proxy.proxy_server import premium_user, prisma_client + from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache try: if not premium_user: @@ -773,6 +780,11 @@ async def delete_project( prisma_models.LiteLLM_ProjectTable | None ) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id}) + await delete_cached_project_object( + project_id=project_id, + user_api_key_cache=user_api_key_cache, + ) + deleted_projects.append(deleted_project) return deleted_projects diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 5489eba1494..a069bd81eca 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.53" +version = "0.1.54" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.53" +version = "0.1.54" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql new file mode 100644 index 00000000000..81b1cbc7ec3 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index beddd899472..fc58ff68b4d 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.83" +version = "0.4.84" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.83" +version = "0.4.84" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm/a2a_protocol/exception_mapping_utils.py b/litellm/a2a_protocol/exception_mapping_utils.py index 16c295f469c..62bef6e02ae 100644 --- a/litellm/a2a_protocol/exception_mapping_utils.py +++ b/litellm/a2a_protocol/exception_mapping_utils.py @@ -202,9 +202,9 @@ async def handle_a2a_localhost_retry( # Fix the agent card URL set_agent_card_url(agent_card, error.base_url) - # Reuse the httpx client LiteLLM attached at creation. It carries this agent's - # trace-id and auth headers, so a fresh client would drop them. Only clients built - # by ``create_a2a_client`` have it; an externally-supplied client cannot be retried. + # Reuse the httpx client and call context LiteLLM attached at creation, since the + # context carries this agent's trace-id/auth headers. Only clients built by + # ``create_a2a_client`` have them; an externally-supplied client cannot be retried. httpx_client: Final = getattr(a2a_client, "_litellm_httpx_client", None) if httpx_client is None: raise RuntimeError( @@ -220,5 +220,8 @@ async def handle_a2a_localhost_retry( ), ) new_client._litellm_httpx_client = httpx_client + new_client._litellm_call_context = getattr( # pyright: ignore[reportAttributeAccessIssue] # LiteLLM-owned stash + a2a_client, "_litellm_call_context", None + ) new_client._litellm_agent_card = agent_card return new_client diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 4b931e84427..322393cd9c4 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -30,6 +30,7 @@ from litellm.utils import client if TYPE_CHECKING: from a2a.client import Client as A2AClientType + from a2a.client import ClientCallContext as A2ACallContextType from a2a.compat.v0_3.types import ( AgentCard, Message, @@ -45,7 +46,7 @@ A2A_SDK_AVAILABLE = False _a2a_conversions: Any = None try: - from a2a.client import Client, ClientConfig, create_client + from a2a.client import Client, ClientCallContext, ClientConfig, create_client from a2a.compat.v0_3 import conversions as _a2a_conversions from a2a.compat.v0_3.types import ( Message, @@ -60,6 +61,7 @@ try: A2A_SDK_AVAILABLE = True except ImportError: Client = None + ClientCallContext = None ClientConfig = None create_client = None @@ -218,6 +220,10 @@ async def _send_message_via_completion_bridge( return LiteLLMSendMessageResponse.from_dict(response_dict, request_id=str(request.id)) +def _get_a2a_call_context(a2a_client: "A2AClientType") -> Optional["A2ACallContextType"]: + return getattr(a2a_client, "_litellm_call_context", None) + + async def _send_message(a2a_client: "A2AClientType", request: "SendMessageRequest") -> "SendMessageResponse": """Send a non-streaming message via a2a-sdk 1.x and return JSON-RPC response.""" if _a2a_conversions is None: @@ -227,7 +233,7 @@ async def _send_message(a2a_client: "A2AClientType", request: "SendMessageReques pb_request: Final = _a2a_conversions.to_core_send_message_request(request) last_event = None - async for event in a2a_client.send_message(pb_request): + async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)): last_event = event if last_event is None: raise RuntimeError("A2A send_message failed: no response received from agent.") @@ -301,7 +307,7 @@ async def _stream_messages( ) pb_request: Final = _a2a_conversions.to_core_send_message_request(request) - async for event in a2a_client.send_message(pb_request): + async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)): compat_chunk = _a2a_conversions.to_compat_stream_response( event, request_id=request.id, @@ -756,26 +762,12 @@ async def create_a2a_client( verbose_logger.info("Creating A2A client for %s", base_url) - # Use get_async_httpx_client with per-agent params so that different agents - # (with different extra_headers) get separate cached clients. The params - # dict is hashed into the cache key, keeping agent auth isolated while - # still reusing connections within the same agent. - # - # Only pass params that AsyncHTTPHandler.__init__ accepts (e.g. timeout). - # Use "disable_aiohttp_transport" key for cache-key-only data (it's - # filtered out before reaching the constructor). - _client_params: Final[dict] = {"timeout": timeout} - if extra_headers: - # Encode headers into a cache-key-only param so each unique header - # set produces a distinct cache key. - _client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items())) _async_handler: Final = get_async_httpx_client( llm_provider=httpxSpecialProvider.A2AProvider, - params=_client_params, + params={"timeout": timeout}, ) httpx_client: Final = _async_handler.client if extra_headers: - httpx_client.headers.update(extra_headers) verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys())) a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall] @@ -784,11 +776,17 @@ async def create_a2a_client( httpx_client=httpx_client, streaming=streaming, ), + resolver_http_kwargs={"headers": extra_headers} if extra_headers else None, ) # Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse - # the configured httpx client (with this agent's trace-id/auth headers) without - # excavating a2a-sdk private internals. + # the configured httpx client and this agent's headers without excavating + # a2a-sdk private internals. a2a_client._litellm_httpx_client = httpx_client + a2a_client._litellm_call_context = ( # pyright: ignore[reportAttributeAccessIssue] # LiteLLM-owned stash + ClientCallContext(service_parameters=extra_headers) # pyright: ignore[reportOptionalCall] # SDK checked above + if extra_headers + else None + ) agent_card: Final = getattr(a2a_client, "_card", None) if agent_card is not None: a2a_client._litellm_agent_card = agent_card diff --git a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py index 339da998d56..024e8c179c2 100644 --- a/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py +++ b/litellm/a2a_protocol/providers/pydantic_ai_agents/transformation.py @@ -6,16 +6,33 @@ This module provides fake streaming by converting non-streaming responses into s """ import asyncio -from collections.abc import AsyncIterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Mapping, Sequence +from typing import Any, Final, Protocol, cast, runtime_checkable from uuid import uuid4 +from pydantic import TypeAdapter + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) +_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_TEXT_ADAPTER: Final = TypeAdapter(str) + + +@runtime_checkable +class _SupportsModelDump(Protocol): + def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ... + + +@runtime_checkable +class _SupportsPydanticDict(Protocol): + def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ... + class PydanticAITransformation: """ @@ -28,7 +45,7 @@ class PydanticAITransformation: """ @staticmethod - def _remove_none_values(obj: Any) -> Any: + def _remove_none_values(obj: object) -> object: """ Recursively remove None values from a dict/list structure. @@ -42,14 +59,18 @@ class PydanticAITransformation: Cleaned object with None values removed """ if isinstance(obj, dict): - return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None} + typed_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(obj) + return {k: PydanticAITransformation._remove_none_values(v) for k, v in typed_dict.items() if v is not None} elif isinstance(obj, list): - return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None] + typed_list: Final = _LIST_ADAPTER.validate_python(obj) + return [PydanticAITransformation._remove_none_values(item) for item in typed_list if item is not None] else: return obj @staticmethod - def _params_to_dict(params: Any) -> dict[str, Any]: + def _params_to_dict( + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", + ) -> Mapping[str, object]: """ Convert params to a dict, handling Pydantic models. @@ -59,10 +80,10 @@ class PydanticAITransformation: Returns: Dict representation of params """ - if hasattr(params, "model_dump"): + if isinstance(params, _SupportsModelDump): # Pydantic v2 model return params.model_dump(mode="python", exclude_none=True) - elif hasattr(params, "dict"): + elif isinstance(params, _SupportsPydanticDict): # Pydantic v1 model return params.dict(exclude_none=True) elif isinstance(params, dict): @@ -75,12 +96,12 @@ class PydanticAITransformation: async def _poll_for_completion( client: AsyncHTTPHandler, endpoint: str, - task_id: str, + task_id: object, request_id: str, max_attempts: int = 30, poll_interval: float = 0.5, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Poll for task completion using tasks/get method. @@ -112,10 +133,10 @@ class PydanticAITransformation: }, ) response.raise_for_status() - poll_data = response.json() + poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) - result = poll_data.get("result", {}) - status = result.get("status", {}) + result = _STR_KEY_DICT_ADAPTER.validate_python(poll_data.get("result", {})) + status = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {})) state = status.get("state", "") verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state) @@ -133,10 +154,10 @@ class PydanticAITransformation: async def _send_and_poll_raw( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -153,14 +174,16 @@ class PydanticAITransformation: Raw Pydantic AI task response (with history/artifacts) """ # Convert params to dict if it's a Pydantic model - params_dict = PydanticAITransformation._params_to_dict(params) - # Remove None values - FastA2A doesn't accept null for optional fields - params_dict = PydanticAITransformation._remove_none_values(params_dict) + params_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python( + PydanticAITransformation._remove_none_values(PydanticAITransformation._params_to_dict(params)) + ) # Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI if "message" in params_dict: - params_dict["message"]["kind"] = "message" + message_value: Final = _ANY_KEY_DICT_ADAPTER.validate_python(params_dict["message"]) + message_value["kind"] = "message" + params_dict["message"] = message_value # Build A2A JSON-RPC request using message/send method for FastA2A compatibility a2a_request: Final = { @@ -189,11 +212,11 @@ class PydanticAITransformation: }, ) response.raise_for_status() - response_data = response.json() + response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json()) # Check if task is already completed - result: Final = response_data.get("result", {}) - status: Final = result.get("status", {}) + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) + status: Final = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {})) state: Final = status.get("state", "") if state != "completed": @@ -217,10 +240,10 @@ class PydanticAITransformation: async def send_non_streaming_request( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a non-streaming A2A request to Pydantic AI agent and wait for completion. @@ -253,10 +276,10 @@ class PydanticAITransformation: async def send_and_get_raw_response( api_base: str, request_id: str, - params: Any, + params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]", timeout: float = 60.0, agent_extra_headers: dict[str, str] | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Send a request to Pydantic AI agent and return the raw task response. @@ -282,9 +305,9 @@ class PydanticAITransformation: @staticmethod def _transform_to_a2a_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform Pydantic AI task response to standard A2A non-streaming format. @@ -328,7 +351,7 @@ class PydanticAITransformation: } @staticmethod - def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]: + def _extract_response_text(response_data: Mapping[str, object]) -> tuple[object, object, Sequence[object]]: """ Extract response text from completed task response. @@ -342,52 +365,53 @@ class PydanticAITransformation: Returns: Tuple of (full_text, message_id, parts) """ - result: Final = response_data.get("result", {}) + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) # Try to extract from artifacts first (preferred for results) artifacts: Final = result.get("artifacts", []) if artifacts: - for artifact in artifacts: - parts = artifact.get("parts", []) + for artifact in _LIST_ADAPTER.validate_python(artifacts): + parts = _LIST_ADAPTER.validate_python(_STR_KEY_DICT_ADAPTER.validate_python(artifact).get("parts", [])) for part in parts: - if part.get("kind") == "text": - text = part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + text = part_dict.get("text", "") if text: return text, str(uuid4()), parts # Fall back to history - get the last agent message - history: Final = result.get("history", []) + history: Final = _LIST_ADAPTER.validate_python(result.get("history", [])) for msg in reversed(history): - if msg.get("role") == "agent": - parts = msg.get("parts", []) - message_id = msg.get("messageId", str(uuid4())) + if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "agent": + parts = _LIST_ADAPTER.validate_python(msg_dict.get("parts", [])) + message_id = msg_dict.get("messageId", str(uuid4())) full_text = "" for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", "")) if full_text: return full_text, message_id, parts # Fall back to message field (original format) message: Final = result.get("message", {}) if message: - parts = message.get("parts", []) - message_id = message.get("messageId", str(uuid4())) + message_dict: Final = _STR_KEY_DICT_ADAPTER.validate_python(message) + parts = _LIST_ADAPTER.validate_python(message_dict.get("parts", [])) + message_id = message_dict.get("messageId", str(uuid4())) full_text = "" for part in parts: - if part.get("kind") == "text": - full_text += part.get("text", "") + if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text": + full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", "")) return full_text, message_id, parts return "", str(uuid4()), [] @staticmethod async def fake_streaming_from_response( - response_data: dict[str, Any], + response_data: Mapping[str, object], request_id: str, chunk_size: int = 50, delay_ms: int = 10, - ) -> AsyncIterator[dict[str, Any]]: + ) -> AsyncIterator[dict[str, object]]: """ Convert a non-streaming A2A response into fake streaming chunks. @@ -410,12 +434,12 @@ class PydanticAITransformation: full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data) # Extract input message from raw response for history - result: Final = response_data.get("result", {}) - history: Final = result.get("history", []) - input_message = {} + result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {})) + history: Final = _LIST_ADAPTER.validate_python(result.get("history", [])) + input_message = _STR_KEY_DICT_ADAPTER.validate_python({}) for msg in history: - if msg.get("role") == "user": - input_message = msg + if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "user": + input_message = msg_dict break # Generate IDs for streaming events @@ -426,45 +450,49 @@ class PydanticAITransformation: # 1. Emit initial task event (kind: "task", status: "submitted") # Format matches A2ACompletionBridgeTransformation.create_task_event - task_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "history": [ - { - "contextId": context_id, - "kind": "message", - "messageId": input_message_id, - "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), - "role": "user", - "taskId": task_id, - } - ], - "id": task_id, - "kind": "task", - "status": { - "state": "submitted", + task_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "history": [ + { + "contextId": context_id, + "kind": "message", + "messageId": input_message_id, + "parts": input_message.get("parts", [{"kind": "text", "text": ""}]), + "role": "user", + "taskId": task_id, + } + ], + "id": task_id, + "kind": "task", + "status": { + "state": "submitted", + }, }, - }, - } + } + ) yield task_event # 2. Emit status update (kind: "status-update", status: "working") # Format matches A2ACompletionBridgeTransformation.create_status_update_event - working_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "final": False, - "kind": "status-update", - "status": { - "state": "working", + working_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": False, + "kind": "status-update", + "status": { + "state": "working", + }, + "taskId": task_id, }, - "taskId": task_id, - }, - } + } + ) yield working_event # Small delay to simulate processing @@ -473,29 +501,32 @@ class PydanticAITransformation: # 3. Emit artifact update chunks (kind: "artifact-update") # Format matches A2ACompletionBridgeTransformation.create_artifact_update_event if full_text: + full_text_str: Final = _TEXT_ADAPTER.validate_python(full_text) # Split text into chunks - for i in range(0, len(full_text), chunk_size): - chunk_text = full_text[i : i + chunk_size] - is_last_chunk = (i + chunk_size) >= len(full_text) + for i in range(0, len(full_text_str), chunk_size): + chunk_text = full_text_str[i : i + chunk_size] + is_last_chunk = (i + chunk_size) >= len(full_text_str) - artifact_event = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "kind": "artifact-update", - "taskId": task_id, - "artifact": { - "artifactId": artifact_id, - "parts": [ - { - "kind": "text", - "text": chunk_text, - } - ], + artifact_event = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "kind": "artifact-update", + "taskId": task_id, + "artifact": { + "artifactId": artifact_id, + "parts": [ + { + "kind": "text", + "text": chunk_text, + } + ], + }, }, - }, - } + } + ) yield artifact_event # Add delay between chunks (except for last chunk) @@ -503,19 +534,21 @@ class PydanticAITransformation: await asyncio.sleep(delay_ms / 1000.0) # 4. Emit final status update (kind: "status-update", status: "completed", final: true) - completed_event: Final = { - "jsonrpc": "2.0", - "id": request_id, - "result": { - "contextId": context_id, - "final": True, - "kind": "status-update", - "status": { - "state": "completed", + completed_event: Final = _STR_KEY_DICT_ADAPTER.validate_python( + { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": True, + "kind": "status-update", + "status": { + "state": "completed", + }, + "taskId": task_id, }, - "taskId": task_id, - }, - } + } + ) yield completed_event verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id) diff --git a/litellm/constants.py b/litellm/constants.py index 75f12190b3e..30d3bb1f26e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_non DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm")) AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview")) ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) +ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) @@ -279,6 +280,7 @@ TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECO GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int( os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) ) +BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000 # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index f23317ae9df..24328549094 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -16,7 +16,10 @@ import asyncio import os import time import traceback +from collections.abc import Mapping +from types import MappingProxyType from typing import Final +from urllib.parse import urlparse from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -27,6 +30,16 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload +DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com" +DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default" + +MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType( + { + "login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE, + "login.microsoftonline.us": "https://monitor.azure.us/.default", + } +) + class AzureSentinelLogger(CustomBatchLogger): """ @@ -42,6 +55,7 @@ class AzureSentinelLogger(CustomBatchLogger): client_id: str | None = None, client_secret: str | None = None, audit_stream_name: str | None = None, + authority_host: str | None = None, **kwargs, ): """ @@ -62,6 +76,10 @@ class AzureSentinelLogger(CustomBatchLogger): If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var. audit_stream_name (str, optional): Stream name from DCR for audit logs. If not provided, will use AZURE_SENTINEL_AUDIT_STREAM_NAME env var or the standard stream name. + authority_host (str, optional): Microsoft Entra authority host that issues the OAuth2 token, + e.g. "https://login.microsoftonline.us" for Azure Government. If not provided, will use + AZURE_SENTINEL_AUTHORITY_HOST or AZURE_AUTHORITY_HOST env vars, or default to the Azure + Public Cloud authority. The Azure Monitor audience is derived from it. """ self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) @@ -76,6 +94,12 @@ class AzureSentinelLogger(CustomBatchLogger): resolved_client_secret: Final = ( client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET") ) + resolved_authority_host: Final = self._normalize_authority_host( + authority_host + or os.getenv("AZURE_SENTINEL_AUTHORITY_HOST") + or os.getenv("AZURE_AUTHORITY_HOST") + or DEFAULT_AZURE_AUTHORITY_HOST + ) if not resolved_dcr_immutable_id: raise ValueError( @@ -119,7 +143,8 @@ class AzureSentinelLogger(CustomBatchLogger): ) # OAuth2 scope for Azure Monitor - self.oauth_scope = "https://monitor.azure.com/.default" + self.authority_host = resolved_authority_host + self.oauth_scope = self._resolve_oauth_scope(authority_host=resolved_authority_host) self.oauth_token: str | None = None self.oauth_token_expires_at: float | None = None @@ -129,6 +154,26 @@ class AzureSentinelLogger(CustomBatchLogger): self.log_queue: list[StandardLoggingPayload] = [] self.audit_log_queue: list[StandardAuditLogPayload] = [] + @staticmethod + def _normalize_authority_host(authority_host: str) -> str: + """ + Normalize an authority host into an absolute URL with no trailing slash. + + Accepts the scheme-qualified form litellm documents ("https://login.microsoftonline.us") + and the bare-host form the azure-identity AzureAuthorityHosts constants use. + """ + stripped: Final = authority_host.strip().rstrip("/") + return stripped if "://" in stripped else f"https://{stripped}" + + @staticmethod + def _resolve_oauth_scope(authority_host: str) -> str: + """ + Map an authority host to the Azure Monitor Logs Ingestion audience for the same cloud, + falling back to the Azure Public Cloud audience for an unrecognized host. + """ + host: Final = urlparse(authority_host).hostname or "" + return MONITOR_SCOPE_BY_AUTHORITY_HOST.get(host, DEFAULT_AZURE_MONITOR_SCOPE) + @staticmethod def _build_api_endpoint(endpoint: str, dcr_immutable_id: str, stream_name: str) -> str: return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01" @@ -150,7 +195,7 @@ class AzureSentinelLogger(CustomBatchLogger): assert self.client_id is not None, "client_id is required" assert self.client_secret is not None, "client_secret is required" - token_url: Final = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" + token_url: Final = f"{self.authority_host}/{self.tenant_id}/oauth2/v2.0/token" token_data: Final = { "client_id": self.client_id, diff --git a/litellm/integrations/otel/__init__.py b/litellm/integrations/otel/__init__.py index 94442e96adb..9c1205bb277 100644 --- a/litellm/integrations/otel/__init__.py +++ b/litellm/integrations/otel/__init__.py @@ -57,6 +57,7 @@ from litellm.integrations.otel.model.semconv import ( Metric, Network, NetworkTransport, + RpcSystem, Server, resolve_operation, resolve_provider, @@ -102,6 +103,7 @@ __all__ = [ "ProxyRequestSpanData", "RequestContext", "RequestIdentity", + "RpcSystem", "Server", "ServerInfo", "ServiceSpanData", diff --git a/litellm/integrations/otel/logger.py b/litellm/integrations/otel/logger.py index e273289168a..2c83406afed 100644 --- a/litellm/integrations/otel/logger.py +++ b/litellm/integrations/otel/logger.py @@ -33,6 +33,7 @@ from litellm.integrations.otel.model.payloads import ( is_mcp_list_tools, is_mcp_tool_call, ) +from litellm.integrations.otel.model.semconv import Error from litellm.integrations.otel.model.spans import SpanRole, span_role_for_service from litellm.integrations.otel.model.utils import to_ns from litellm.integrations.otel.plumbing.context import ( @@ -634,18 +635,23 @@ class OpenTelemetryV2(CustomLogger): """Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a failure that dies before any LLM-call span exists (malformed body, auth / validation rejection). Called from the proxy's global exception handler via - ``_close_dangling_otel_server_span``. The instrumentor still owns the span's - status and lifecycle, so this only decorates it — never sets status, never - ends it — and emits no exception event, matching v1's SERVER-span behavior - and avoiding a duplicate of the event ``async_post_call_failure_hook`` or - the ``auth`` phase span already records.""" + ``_close_dangling_otel_server_span``, which swallows the exception into a + ``JSONResponse`` so the instrumentor never sees it and leaves the span + ``UNSET``; the status is set here instead (v1 did the same from the handler) + so a failed request reads as failed and not merely as a span carrying an + error message. The instrumentor still owns the span's lifecycle, so this + never ends it. The exception event is recorded only when nothing stamped + this span already — ``async_post_call_failure_hook`` and the ``auth`` phase + span record their own, and a second event would duplicate it — while the + attributes are always restamped so ``error.code`` stays pinned to the real + response status.""" if span is None or not is_recordable_span(span): return + already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ()) stamp_error( span, _span_error_from_exception(exception, status_code=status_code), - record_event=False, - set_status=False, + record_event=not already_stamped, ) async def async_post_call_failure_hook( diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index af56734bec1..032441535e0 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -31,7 +31,9 @@ from litellm.integrations.otel.model.semconv import ( MCP, Error, GenAI, + JsonRpc, LiteLLM, + RpcSystem, Server, ) from litellm.integrations.otel.model.spans import db_system @@ -94,11 +96,14 @@ class GenAIMapper: _MCP_ATTRS: dict[str, Callable[[MCPToolCallSpanData], AttrValue | None]] = { GenAI.OPERATION_NAME: lambda d: d.operation.value, + JsonRpc.SYSTEM: lambda d: RpcSystem.JSONRPC.value if d.server_address and d.server_port else None, MCP.METHOD_NAME: lambda d: d.method, MCP.SESSION_ID: lambda d: d.session_id, GenAI.TOOL_NAME: lambda d: d.tool_name or None, GenAI.TOOL_CALL_ARGUMENTS: lambda d: d.arguments_json, GenAI.TOOL_CALL_RESULT: lambda d: d.result_json, + Server.ADDRESS: lambda d: d.server_address, + Server.PORT: lambda d: d.server_port, LiteLLM.MCP_SERVER_NAME: lambda d: d.server_name, LiteLLM.CALL_ID: lambda d: d.identity.call_id or None, f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost, diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index 02499010624..aba9cc80240 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -364,6 +364,34 @@ class LLMCallSpanData: # --- the MCP tool-call model ------------------------------------------------- # +def _upstream_address_port(resource: str | None) -> tuple[str | None, int | None]: + """Split a redacted MCP server origin into ``server.address`` / ``server.port``. + + ``mcp_server_resource`` is a scheme + host + port origin with userinfo, path, + query and fragment already stripped. The port falls back to the scheme default + when the origin omits it, because a consumer that keys a downstream dependency + off the address renders a missing port as ``0``. + + The origin is rebuilt without its IPv6 brackets upstream, so reading the port can + raise on an address the host check still admits: a zone-scoped ``fe80::1%25eth0`` + leaves a truthy hostname of ``fe80`` behind. Both halves are read inside the guard + so an unparseable origin yields no address rather than propagating out of span + construction, matching how the redactor guards the same split. + """ + if not resource: + return None, None + try: + parsed: Final = urlsplit(resource) + hostname: Final = parsed.hostname + port: Final = parsed.port + except ValueError: + return None, None + if not hostname: + return None, None + default_port: Final = 443 if parsed.scheme == "https" else 80 if parsed.scheme == "http" else None + return hostname, port or default_port + + @dataclass(frozen=True) class MCPToolCallSpanData: """One MCP ``tools/call`` execution, parsed from a closed request's payload. @@ -378,6 +406,8 @@ class MCPToolCallSpanData: method: str tool_name: str server_name: str | None + server_address: str | None + server_port: int | None session_id: str | None arguments_json: str | None result_json: str | None @@ -390,11 +420,14 @@ class MCPToolCallSpanData: cls, payload: StandardLoggingPayload, capture_content: bool = False ) -> MCPToolCallSpanData: meta: Final = _mcp_tool_call_metadata(cast(Mapping[str, object], payload)) + address, port = _upstream_address_port(as_str(meta.get("mcp_server_resource")) or None) return cls( operation=resolve_operation(as_str(payload.get("call_type"))), method=MCPMethod.TOOLS_CALL.value, tool_name=as_str(meta.get("name")) or "", server_name=as_str(meta.get("mcp_server_name")), + server_address=address, + server_port=port, session_id=as_str(meta.get("mcp_session_id")), arguments_json=( _json_or_none(meta.get("arguments")) if capture_content and meta.get("arguments") is not None else None diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 24f0b947b08..3d585c36b67 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -130,11 +130,25 @@ class JsonRpc: """JSON-RPC keys carried on MCP spans. The error/status code lives in the ``rpc.*`` namespace per semconv, not ``jsonrpc.*``.""" + SYSTEM: Final = "rpc.system" REQUEST_ID: Final = "jsonrpc.request.id" PROTOCOL_VERSION: Final = "jsonrpc.protocol.version" RESPONSE_STATUS_CODE: Final = "rpc.response.status_code" +class RpcSystem(str, Enum): + """Well-known values for ``rpc.system``. MCP frames every message as JSON-RPC 2.0. + + Naming the system also classifies the span: a CLIENT span carrying none of the + ``rpc.*``/``http.*``/``db.*``/``messaging.*`` families records no span type or + subtype in backends that derive those from the attribute family. It is emitted + only alongside ``server.address``/``server.port``, since a backend that reads it + as a downstream dependency names that dependency from the server address. + """ + + JSONRPC = "jsonrpc" + + class NetworkTransport(str, Enum): """Well-known values for ``network.transport``.""" diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 71388134e98..2c5f7484dac 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -9,7 +9,7 @@ server-side using litellm router's search tools. import asyncio import math import uuid -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -37,10 +37,17 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, ) from litellm.types.integrations.websearch_interception import ( + AnthropicSearchQuery, + AnthropicServerToolUseBlock, WebSearchInterceptionConfig, ) from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import CallTypes, LlmProviders +from litellm.types.utils import ( + AgenticLoopParams, + CallTypes, + LlmProviders, + StandardLoggingUserAPIKeyMetadata, +) from litellm.utils import ProviderConfigManager if TYPE_CHECKING: @@ -263,7 +270,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None # Check if request has tools with native web_search - tools: Final = kwargs.get("tools") + tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools") if not tools: return None @@ -312,7 +319,9 @@ class WebSearchInterceptionLogger(CustomLogger): return kwargs - def _convert_responses_tools(self, kwargs: Mapping[str, object], tools: list[dict[str, object]]) -> dict | None: + def _convert_responses_tools( + self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]] + ) -> dict[str, object] | None: """Convert Responses API web search tools to the LiteLLM standard function tool.""" if not any(is_web_search_tool_responses(tool) for tool in tools): return None @@ -377,7 +386,7 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _tool_name(tool: dict[str, Any]) -> str | None: + def _tool_name(tool: Mapping[str, object]) -> object: """Effective tool name, handling OpenAI ``function`` wrapper shape.""" fn: Final = tool.get("function") if tool.get("type") == "function" and isinstance(fn, dict): @@ -833,22 +842,48 @@ class WebSearchInterceptionLogger(CustomLogger): def _build_native_result_blocks( tool_calls: list[dict], structured_results: list[SearchResponse | None], - ) -> list[dict[str, object]]: - """Build one ``web_search_tool_result`` block per tool_call.""" - blocks: Final[list[dict[str, object]]] = [] - for i, tool_call in enumerate(tool_calls): - tool_use_id = tool_call.get("id") or "" - structured = structured_results[i] if i < len(structured_results) else None - blocks.append( - WebSearchTransformation.build_web_search_tool_result_block( - tool_use_id=tool_use_id, - search_response=structured, - ) + ) -> tuple[Mapping[str, object], ...]: + """ + Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call. + + The pair is what Anthropic's spec requires: a bare result block, or one + keyed by the model's ``toolu_...`` id instead of a ``srvtoolu_...`` one, + is rejected on replay ("String should match pattern '^srvtoolu_'") and + leaves native clients without a search to attach the sources to. + """ + return tuple( + block + for i, tool_call in enumerate(tool_calls) + for block in WebSearchInterceptionLogger._native_result_pair( + query=WebSearchInterceptionLogger._tool_call_query(tool_call), + search_response=structured_results[i] if i < len(structured_results) else None, ) - return blocks + ) @staticmethod - def _inject_native_blocks(response: Any, native_blocks: list[dict[str, object]]) -> Any: + def _tool_call_query(tool_call: Mapping[str, object]) -> str: + tool_input: Final = tool_call.get("input") + if not isinstance(tool_input, Mapping): + return "" + query: Final = tool_input.get("query") + return query if isinstance(query, str) else "" + + @staticmethod + def _native_result_pair( + query: str, + search_response: SearchResponse | None, + ) -> tuple[Mapping[str, object], Mapping[str, object]]: + tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}" + return ( + AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(), + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id=tool_use_id, + search_response=search_response, + ), + ) + + @staticmethod + def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response @@ -1243,7 +1278,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs) if logging_obj is not None: - agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {}) + agentic_params: Final[AgenticLoopParams] = logging_obj.model_call_details.get("agentic_loop_params", {}) full_model_name = agentic_params.get("model", model) verbose_logger.debug( "WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]", @@ -1288,6 +1323,7 @@ class WebSearchInterceptionLogger(CustomLogger): search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None search_litellm_params: dict[str, Any] = {} + search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) if search_tool is not None: await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) search_litellm_params = dict(search_tool.get("litellm_params", {}) or {}) @@ -1304,12 +1340,30 @@ class WebSearchInterceptionLogger(CustomLogger): verbose_logger.debug( "WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider ) + user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs) + search_metadata: Final = ( + None + if user_api_key_auth is None + else self._build_search_request_metadata( + user_api_key_auth=user_api_key_auth, + search_tool_name=search_tool_name, + ) + ) search_kwargs: Final = { key: value for key, value in search_litellm_params.items() if key != "search_provider" and value is not None } - result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + result: Final = ( + await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + if search_metadata is None + else await litellm.asearch( + query=query, + search_provider=search_provider, + litellm_metadata=search_metadata, + **search_kwargs, + ) + ) # Format using transformation function search_result_text: Final = WebSearchTransformation.format_search_response(result) @@ -1366,6 +1420,35 @@ class WebSearchInterceptionLogger(CustomLogger): team_object=team_object, ) + @staticmethod + def _build_search_request_metadata( + user_api_key_auth: "UserAPIKeyAuth", + search_tool_name: str | None, + ) -> Mapping[str, object]: + """ + Spend-tracking metadata for the intercepted search, so its provider cost is logged + and billed against the key/user/team that made the originating LLM request instead + of being dropped by the proxy's spend hook for lack of an owner. + """ + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = ( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth) + ) + return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches + **user_api_key_metadata, + "model_group": search_tool_name, + "user_api_key": user_api_key_auth.api_key, + "user_api_key_auth": user_api_key_auth, + } + + @staticmethod + def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None: + if search_tool is None: + return None + search_tool_name: Final = search_tool.get("search_tool_name") + return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None + @staticmethod def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None": if not kwargs: diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 795810a7c40..199ab020559 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -412,6 +412,15 @@ class WebSearchTransformation: block that should accompany the model's text reply when the original request used a native ``web_search_*`` tool. + The spec'd shape carries page text only in ``encrypted_content``, an + opaque server-issued blob that we cannot mint. Emitting the four spec + fields alone would drop the snippet entirely, leaving the client (and + the model, on any replayed follow-up turn) with URLs and titles but no + evidence to answer from, forcing a fetch per result. So the snippet is + carried in an additive ``snippet`` key alongside the spec fields. + ``encrypted_content`` stays empty rather than holding plaintext, which + would assert encryption semantics that do not hold. + Spec reference: https://docs.anthropic.com/en/api/web-search-tool @@ -438,6 +447,7 @@ class WebSearchTransformation: "title": title, "page_age": page_age, "encrypted_content": "", + "snippet": getattr(r, "snippet", "") or "", } ) return { diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index d6433ad3332..f251ab4d74a 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -115,8 +115,11 @@ def get_litellm_params( litellm_request_debug: bool | None = None, **kwargs, ) -> dict: + _litellm_metadata_dict: Final = litellm_metadata if isinstance(litellm_metadata, dict) else None + resolved_metadata: Final = _litellm_metadata_dict.copy() if not metadata and _litellm_metadata_dict else metadata + # Derive litellm_session_id / litellm_trace_id from metadata when not provided (call chaining) - _meta: Final = metadata or {} + _meta: Final = resolved_metadata or {} if litellm_session_id is None: litellm_session_id = _meta.get("session_id") or _meta.get("trace_id") if litellm_trace_id is None: @@ -139,7 +142,7 @@ def get_litellm_params( "model_alias_map": model_alias_map, "completion_call_id": completion_call_id, "aembedding": aembedding, - "metadata": metadata, + "metadata": resolved_metadata, "model_info": model_info, "proxy_server_request": proxy_server_request, "preset_cache_key": preset_cache_key, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9475441e214..99721c3ffa2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -585,8 +585,8 @@ class Logging(LiteLLMLoggingBaseClass): """ base_litellm_params: Final[dict[str, Any]] = {} - if "metadata" in kwargs: - base_litellm_params["metadata"] = kwargs["metadata"] + if isinstance(kwargs.get("metadata"), dict): + base_litellm_params["metadata"] = kwargs["metadata"].copy() if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict): base_litellm_params["litellm_metadata"] = kwargs["litellm_metadata"] if "metadata" not in base_litellm_params: diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 3f43fe38f5e..c596e821ce9 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -6,7 +6,7 @@ import io import json import mimetypes import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from os import PathLike from pathlib import Path from typing import TYPE_CHECKING, Any, Final, Literal, cast @@ -1775,3 +1775,24 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]: idx = end_idx return results + + +def text_completion_prompt_to_messages(prompt: object) -> tuple[AllMessageValues, ...]: + """ + Wrap an OpenAI ``/v1/completions`` ``prompt`` into Chat Completion messages. + + Mirrors what ``litellm.text_completion`` does on the real-time path: a + string becomes a single user message, and a list of strings becomes one + user message per element. Pre-tokenized prompts (``list[int]`` / + ``list[list[int]]``) are only meaningful for the OpenAI-family text + endpoints, so they are rejected here rather than silently forwarded, as is + an empty prompt, which every chat-shaped provider rejects downstream. + """ + prompt_type_name: Final = type(prompt).__name__ + if isinstance(prompt, str) and prompt: + return (ChatCompletionUserMessage(role="user", content=prompt),) + entries: Final = cast("Sequence[object]", prompt) if isinstance(prompt, Sequence) else () + string_entries: Final = tuple(entry for entry in entries if isinstance(entry, str) and entry) + if string_entries and len(string_entries) == len(entries): + return tuple(ChatCompletionUserMessage(role="user", content=entry) for entry in string_entries) + raise ValueError(f"`prompt` must be a non-empty string or a non-empty list of strings. Got: {prompt_type_name}.") diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 4ab109e5830..4b774269cc2 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,5 +1,6 @@ import asyncio import json +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm @@ -21,12 +22,24 @@ if TYPE_CHECKING: from websockets.asyncio.client import ClientConnection from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.guardrails import GuardrailEventHooks CLIENT_CONNECTION_CLASS = ClientConnection else: CLIENT_CONNECTION_CLASS = Any +class _ClientWebSocketExceptions(Protocol): + ConnectionClosed: type[Exception] + + +class _ClientWebSocket(Protocol): + exceptions: _ClientWebSocketExceptions + + async def send_text(self, data: str) -> None: ... + async def receive_text(self) -> str: ... + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... def normalize(self, event: dict) -> dict: ... @@ -56,7 +69,7 @@ class RealTimeStreaming: force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, ): - self.websocket = websocket + self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.messages: list[OpenAIRealtimeEvents] = [] @@ -85,7 +98,7 @@ class RealTimeStreaming: self.current_item_chunks: list[OpenAIRealtimeOutputItemDone] | None = None self.current_delta_type: ALL_DELTA_TYPES | None = None self.session_configuration_request: str | None = None - self.user_api_key_dict: "UserAPIKeyAuth | None" = user_api_key_dict + self.user_api_key_dict = user_api_key_dict self.request_data: dict = request_data or {} # Violation counter for end_session_after_n_fails support self._violation_count: int = 0 @@ -129,7 +142,7 @@ class RealTimeStreaming: ] ) _CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"]) - _AUDIO_FORMAT_MAP: dict[str, dict[str, Any]] = { + _AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = { "pcm16": {"type": "audio/pcm", "rate": 24000}, "g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000}, "g711_alaw": {"type": "audio/G711-alaw", "rate": 8000}, @@ -283,6 +296,7 @@ class RealTimeStreaming: if event_obj.get("type") != "response.done": return response: Final = cast(dict[str, Any], event_obj.get("response", {})) + item: Mapping[str, object] for item in response.get("output", []): if item.get("type") == "function_call": self.tool_calls.append( @@ -386,7 +400,7 @@ class RealTimeStreaming: return message try: - message_obj: Final = json.loads(message) + message_obj: Final[Mapping[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return message @@ -489,7 +503,7 @@ class RealTimeStreaming: if self._backend_setup_complete and not self._flushing_pending_messages_until_setup: return False try: - msg_obj: Final = json.loads(message) + msg_obj: Final[Mapping[str, object]] = json.loads(message) except (json.JSONDecodeError, TypeError): return False return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES @@ -557,7 +571,7 @@ class RealTimeStreaming: def _event_to_client_json(self, event: dict) -> str: return json.dumps(self._normalize_event_for_ga_client(event)) - async def _send_event_to_client(self, event: Any, event_str: str) -> bool: + async def _send_event_to_client(self, event: object, event_str: str) -> bool: if self._should_drop_event_from_client(event): return False if isinstance(event, dict): @@ -597,12 +611,12 @@ class RealTimeStreaming: def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" - turn_detection: Final[dict[str, Any]] = { + turn_detection: Final[dict[str, str | bool]] = { "type": "server_vad", "create_response": False, } if self._backend_uses_beta_protocol: - session: dict[str, Any] = {"turn_detection": turn_detection} + session: dict[str, object] = {"turn_detection": turn_detection} else: session = { "type": "realtime", @@ -656,7 +670,7 @@ class RealTimeStreaming: def _has_realtime_guardrails_for_event_hooks( self, - event_hooks: list[Any], + event_hooks: Sequence["GuardrailEventHooks"], ) -> bool: """Return True if any callback would run for one of ``event_hooks``.""" from litellm.integrations.custom_guardrail import CustomGuardrail @@ -701,7 +715,7 @@ class RealTimeStreaming: transcript: str, item_id: str | None = None, pre_block_backend_message: str | None = None, - event_hooks: list[Any] | None = None, + event_hooks: Sequence["GuardrailEventHooks"] | None = None, ) -> bool: """ Run registered guardrails on realtime text (transcript, user message, tool output). @@ -755,7 +769,7 @@ class RealTimeStreaming: raise # Extract the human-readable error from the detail dict (HTTPException) # or fall back to str(e) for plain ValueError. - detail = getattr(e, "detail", None) + detail: object | None = getattr(e, "detail", None) if isinstance(detail, dict): safe_msg = detail.get("error") or str(e) elif detail is not None: @@ -828,7 +842,7 @@ class RealTimeStreaming: return True return False - async def _handle_provider_config_message(self, raw_response) -> None: + async def _handle_provider_config_message(self, raw_response: str) -> None: """Process a backend message when a provider_config is set (transformed path).""" returned_object: Final = self.provider_config.transform_realtime_response( raw_response, @@ -912,7 +926,7 @@ class RealTimeStreaming: await self._send_event_to_client(event, event_str) @staticmethod - def _parse_backend_event(raw_response: str) -> dict | None: + def _parse_backend_event(raw_response: str) -> dict[str, object] | None: """Parse a backend frame once. Returns None for non-JSON or non-object frames.""" try: event: Final = json.loads(raw_response) @@ -1022,7 +1036,7 @@ class RealTimeStreaming: objects and any test doubles that expose a .scope dict. """ try: - headers: Final = websocket.scope.get("headers", []) + headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", []) for name, value in headers: if isinstance(name, bytes): name = name.decode("latin-1") @@ -1073,9 +1087,9 @@ class RealTimeStreaming: session["output_modalities"] = ["text"] # 3-7. Lift flat audio fields into the nested audio object - audio: Final[dict[str, Any]] = {} - inp: Final[dict[str, Any]] = {} - out: Final[dict[str, Any]] = {} + audio: Final[dict[str, object]] = {} + inp: Final[dict[str, object]] = {} + out: Final[dict[str, object]] = {} # voice → audio.output.voice if "voice" in session: @@ -1192,7 +1206,7 @@ class RealTimeStreaming: # model; check them with the same guardrail used for # user text so an attacker cannot smuggle blocked # content into a function_call_output. - output = item.get("output", "") + output: object = item.get("output", "") output_text = output if isinstance(output, str) else json.dumps(output) if output_text: # Build the sanitized function_call_output up @@ -1243,7 +1257,7 @@ class RealTimeStreaming: # interaction turn. continue elif item.get("role") == "user": - content_list = item.get("content", []) + content_list: Sequence[object] = item.get("content", []) texts = [ c.get("text", "") for c in content_list @@ -1282,7 +1296,7 @@ class RealTimeStreaming: and not self._guardrail_turn_detection_update_sent and self._has_audio_transcription_guardrails() ): - session = msg_obj.setdefault("session", {}) + session: object = msg_obj.setdefault("session", {}) if isinstance(session, dict): existing_td = session.get("turn_detection") if not isinstance(existing_td, dict): diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 3f967e29002..886ba6a3a18 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -3,7 +3,7 @@ import time from collections.abc import Iterator, Mapping, Sequence from itertools import groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Union, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, Union, cast from litellm._logging import verbose_logger from litellm.types.llms.openai import ( @@ -30,6 +30,7 @@ from litellm.types.utils import ( from litellm.utils import print_verbose, token_counter if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, ) @@ -39,6 +40,60 @@ if TYPE_CHECKING: ) +class _ThinkingBlockFragment(TypedDict, total=False): + type: str | None + data: str | None + thinking: str | None + signature: str | None + + +class _ThinkingDelta(TypedDict, total=False): + thinking_blocks: Sequence[_ThinkingBlockFragment] + + +class _ThinkingChoice(TypedDict, total=False): + delta: _ThinkingDelta + + +class _ThinkingChunk(TypedDict): + choices: Sequence[_ThinkingChoice] + + +class _ContentChoice(TypedDict, total=False): + delta: Mapping[str, str | None] + + +class _ContentChunk(TypedDict): + choices: Sequence[_ContentChoice] + + +class _AudioDelta(TypedDict, total=False): + audio: ChatCompletionAudioDelta | None + + +class _AudioChoice(TypedDict, total=False): + delta: _AudioDelta + + +class _AudioChunk(TypedDict): + choices: Sequence[_AudioChoice] + + +class _UsageBearingChunk(TypedDict, total=False): + usage: Usage | None + _hidden_params: Mapping[str, str] + + +class _UsageSummary(TypedDict): + prompt_tokens: int | None + completion_tokens: int | None + cache_creation_input_tokens: int | None + cache_read_input_tokens: int | None + completion_tokens_details: CompletionTokensDetails | None + prompt_tokens_details: PromptTokensDetailsWrapper | None + cost: float | None + + def capture_cache_creation_token_details( prompt_tokens_details: PromptTokensDetailsWrapper | None, current: CacheCreationTokenDetails | None, @@ -78,7 +133,7 @@ class ChunkProcessor: return [] first_chunk: Final = chunks[0] - first_hidden_params: dict[str, Any] = {} + first_hidden_params: dict[str, object] = {} if isinstance(first_chunk, dict): candidate = first_chunk.get("_hidden_params", {}) if isinstance(candidate, dict): @@ -115,8 +170,8 @@ class ChunkProcessor: @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, - chunks: list[Any], - logging_obj: Any | None = None, + chunks: list[object], + logging_obj: "Logging | None" = None, ) -> None: if not chunks: return @@ -456,7 +511,7 @@ class ChunkProcessor: ) def get_combined_content( - self, chunks: list[dict[str, Any]], delta_key: str = "content" + self, chunks: Sequence["_ContentChunk"], delta_key: str = "content" ) -> ChatCompletionAssistantContentValue: content_list: Final[list[str]] = [] for chunk in chunks: @@ -475,7 +530,7 @@ class ChunkProcessor: return combined_content def get_combined_thinking_content( - self, chunks: list[dict[str, Any]] + self, chunks: Sequence["_ThinkingChunk"] ) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None: from litellm.types.llms.openai import ( ChatCompletionRedactedThinkingBlock, @@ -532,10 +587,10 @@ class ChunkProcessor: return thinking_blocks return None - def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue: + def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue: return self.get_combined_content(chunks, delta_key="reasoning_content") - def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse: + def get_combined_audio_content(self, chunks: Sequence["_AudioChunk"]) -> ChatCompletionAudioResponse: base64_data_list: Final[list[str]] = [] transcript_list: Final[list[str]] = [] expires_at: int | None = None @@ -544,7 +599,7 @@ class ChunkProcessor: for chunk in chunks: choices = chunk["choices"] for choice in choices: - delta = choice.get("delta") or {} + delta: _AudioDelta = choice.get("delta") or {} audio: ChatCompletionAudioDelta | None = delta.get("audio") if audio is not None: for k, v in audio.items(): @@ -565,7 +620,7 @@ class ChunkProcessor: id=id, ) - def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict: + def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> "_UsageSummary": prompt_tokens = 0 completion_tokens = 0 ## anthropic prompt caching information ## @@ -623,8 +678,8 @@ class ChunkProcessor: return reasoning_tokens @staticmethod - def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None: - usage_chunk: Usage | dict[str, Any] | None = None + def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None: + usage_chunk: Usage | None = None if hasattr(chunk, "usage") and chunk.usage is not None: usage_chunk = chunk.usage elif "usage" in chunk: @@ -640,7 +695,7 @@ class ChunkProcessor: def _calculate_usage_per_chunk( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], ) -> "UsagePerChunk": from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import ( UsagePerChunk, @@ -721,13 +776,7 @@ class ChunkProcessor: "web_search_requests", ) - prompt_tokens_details = ( - cast( - PromptTokensDetailsWrapper | None, - usage_chunk_dict["prompt_tokens_details"], - ) - or prompt_tokens_details - ) + prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details cache_creation_token_details = capture_cache_creation_token_details( prompt_tokens_details, cache_creation_token_details @@ -758,7 +807,7 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], completion_tokens: int, completion_usage_updates: int, ) -> int: @@ -797,7 +846,7 @@ class ChunkProcessor: def calculate_usage( self, - chunks: list[dict[str, Any] | ModelResponse], + chunks: Sequence["_UsageBearingChunk | ModelResponse"], model: str, completion_output: str, messages: list | None = None, @@ -851,8 +900,8 @@ class ChunkProcessor: setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens) # for anthropic if completion_tokens_details is not None: if isinstance(completion_tokens_details, CompletionTokensDetails): - returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper( - **completion_tokens_details.model_dump() + returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper.model_validate( + completion_tokens_details.model_dump() ) else: returned_usage.completion_tokens_details = completion_tokens_details diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 314cfef6d84..9aa5a4f465f 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -4,9 +4,12 @@ This file contains common utils for anthropic calls. import copy import re -from typing import Any, Final +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Any, Final, Literal import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -1057,6 +1060,152 @@ def sanitize_tool_use_ids_in_anthropic_messages(messages: list[Any]) -> list[Any return out +class _ReplayedSearchQuery(BaseModel): + model_config = ConfigDict(extra="allow") + + query: str = "" + + +class _ReplayedWebSearchResult(BaseModel): + model_config = ConfigDict(extra="allow") + + type: Literal["web_search_result"] + url: str = "" + title: str = "" + snippet: str = "" + encrypted_content: str = "" + + +class _ReplayedWebSearchToolResult(BaseModel): + model_config = ConfigDict(extra="allow") + + type: Literal["web_search_tool_result"] + tool_use_id: str + content: tuple[_ReplayedWebSearchResult, ...] + + +class _ReplayedServerToolUse(BaseModel): + model_config = ConfigDict(extra="allow") + + type: Literal["server_tool_use"] + id: str + input: _ReplayedSearchQuery = _ReplayedSearchQuery() + + +class _TextBlock(BaseModel): + type: Literal["text"] = "text" + text: str + + +_WEB_SEARCH_TOOL_RESULT_ADAPTER: Final = TypeAdapter(_ReplayedWebSearchToolResult) +_SERVER_TOOL_USE_ADAPTER: Final = TypeAdapter(_ReplayedServerToolUse) + + +def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchToolResult | None: + """ + The parsed block when it is a ``web_search_tool_result`` carrying no + ``encrypted_content``, else None for anything Anthropic itself issued. + + An empty ``content`` list is flattenable too. It is what the interceptor emits + when a search legitimately returns nothing and when a search raises, and it + carries neither evidence to preserve nor an ``encrypted_content`` to respect, + so leaving it in place only buys the 400 this whole function exists to avoid. + """ + try: + parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block) + except ValidationError: + return None + if any(result.encrypted_content for result in parsed.content): + return None + return parsed + + +def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None: + try: + return _SERVER_TOOL_USE_ADAPTER.validate_python(block) + except ValidationError: + return None + + +def _render_web_search_results(query: str, results: tuple[_ReplayedWebSearchResult, ...]) -> str: + header: Final = f"Web search results for '{query}':" if query else "Web search results:" + if not results: + return f"{header}\n\nNo results were returned." + body: Final = "\n\n".join( + "\n".join( + line + for line in ( + f"Title: {result.title}" if result.title else "", + f"URL: {result.url}" if result.url else "", + f"Snippet: {result.snippet}" if result.snippet else "", + ) + if line + ) + for result in results + ) + return f"{header}\n\n{body}" if body else header + + +def _rewrite_replayed_web_search_block( + block: object, + flattenable: Mapping[str, _ReplayedWebSearchToolResult], + queries: Mapping[str, str], +) -> object | None: + parsed_result: Final = _flattenable_web_search_tool_result(block) + if parsed_result is not None: + return _TextBlock( + text=_render_web_search_results(queries.get(parsed_result.tool_use_id, ""), parsed_result.content) + ).model_dump() + parsed_use: Final = _replayed_server_tool_use(block) + if parsed_use is not None and parsed_use.id in flattenable: + return None + return block + + +def _flatten_web_search_results_in_message(message: object) -> object: + if not isinstance(message, Mapping) or not isinstance(message.get("content"), Sequence): + return message + content: Final = message["content"] + if isinstance(content, str): + return message + flattenable: Final = MappingProxyType( + { + parsed.tool_use_id: parsed + for parsed in (_flattenable_web_search_tool_result(block) for block in content) + if parsed is not None + } + ) + if not flattenable: + return message + queries: Final = MappingProxyType( + { + parsed.id: parsed.input.query + for parsed in (_replayed_server_tool_use(block) for block in content) + if parsed is not None + } + ) + rewritten: Final = tuple(_rewrite_replayed_web_search_block(block, flattenable, queries) for block in content) + return {**message, "content": [b for b in rewritten if b is not None]} # mutable-ok: JSON wire format + + +def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: as sibling sanitizers + messages: list[Any], +) -> list[Any]: + """ + Return a new message list with replayed ``web_search_tool_result`` blocks that + carry no ``encrypted_content`` rewritten into plain ``text`` blocks holding the + same title / url / snippet evidence. + + ``encrypted_content`` is an opaque blob only Anthropic's own search backend can + mint, so blocks synthesized by LiteLLM (websearch interception against a search + provider) are rejected with ``Invalid encrypted_content in search_result block`` + when a native client loops them back as history. Flattening them keeps the + evidence in the conversation instead of 400ing the follow-up turn, and leaves + genuine Anthropic-issued blocks untouched. + """ + return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format + + def process_anthropic_headers(headers: httpx.Headers | dict) -> dict: openai_headers: Final = {} if "anthropic-ratelimit-requests-limit" in headers: diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py index a9751489473..36f3e875a7e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/handler.py @@ -1,8 +1,9 @@ -from collections.abc import AsyncIterator, Coroutine, Iterator +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from typing import ( TYPE_CHECKING, Any, Final, + TypeAlias, cast, ) @@ -33,8 +34,12 @@ if TYPE_CHECKING: # Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge. ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config"}) +_AnthropicMessages: TypeAlias = "list[dict[str, object]]" +_AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None" +_ContextManagementSpec: TypeAlias = "dict[str, object] | list[dict[str, object]] | None" -def _messages_have_compaction_block(messages: list[dict]) -> bool: + +def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: content = msg.get("content") @@ -54,8 +59,10 @@ def _proxy_router_fallback() -> "Router | None": return _proxy_router -def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None: - """Return ``kwargs["litellm_metadata"]`` when it's a dict; ``None`` otherwise. +def _extract_proxy_litellm_metadata( + kwargs: Mapping[str, object], +) -> "tuple[dict[str, object], UserAPIKeyAuth | None] | tuple[None, None]": + """Return ``(kwargs["litellm_metadata"], its user_api_key_auth)`` when it's a dict; ``(None, None)`` otherwise. The proxy attaches its auth/spend-attribution fields (``user_api_key``, ``user_api_key_team_id``, ``litellm_call_id``, the full ``UserAPIKeyAuth`` @@ -68,18 +75,19 @@ def _extract_proxy_litellm_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | """ litellm_metadata: Final = kwargs.get("litellm_metadata") if not isinstance(litellm_metadata, dict): - return None - return litellm_metadata + return None, None + user_api_key_auth: Final[UserAPIKeyAuth | None] = litellm_metadata.get("user_api_key_auth") + return litellm_metadata, user_api_key_auth async def _prepare_context_managed_request( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, + messages: _AnthropicMessages, + tools: list[dict[str, object]] | None, + system: _AnthropicSystem, + context_management_spec: _ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, llm_router: "Router | None", user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -102,11 +110,11 @@ async def _prepare_context_managed_request( if polyfill_will_run: history_result: PolyfillResult | None = None - working_messages: list[dict] = messages - working_system: Any | None = system + working_messages: _AnthropicMessages = messages + working_system: _AnthropicSystem = system else: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) working_messages = history_result.messages if history_result is not None else messages @@ -136,7 +144,7 @@ async def _prepare_context_managed_request( # to non-Anthropic backends that would reject them. if polyfill_will_run and history_result is None: history_result = apply_client_compaction_block_history( - messages=cast(list[dict[str, Any]], messages), + messages=messages, system=system, ) return history_result @@ -144,7 +152,7 @@ async def _prepare_context_managed_request( def _polyfill_will_run( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when ``compact_20260112`` will run via the polyfill dispatcher. @@ -171,7 +179,7 @@ def _polyfill_will_run( def _spec_has_non_compact_edits( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, ) -> bool: """Return True when the spec includes edits other than ``compact_20260112``. @@ -209,9 +217,9 @@ def _context_management_explicitly_dropped(additional_drop_params: list[str] | N def _normalize_spec_edits( *, - context_management_spec: Any, + context_management_spec: _ContextManagementSpec, additional_drop_params: list[str] | None, -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """Return the normalized ``edits`` list, or ``None`` if the polyfill won't run. Delegates spec-shape normalization to the dispatcher's ``_normalize_spec`` @@ -236,11 +244,11 @@ def _normalize_spec_edits( async def _run_polyfill_if_enabled( *, model: str, - messages: list[dict], - tools: list[dict] | None, - system: Any | None, - context_management_spec: Any, - litellm_metadata: dict | None, + messages: _AnthropicMessages, + tools: list[dict[str, object]] | None, + system: _AnthropicSystem, + context_management_spec: _ContextManagementSpec, + litellm_metadata: dict[str, object] | None, additional_drop_params: list[str] | None, llm_router: "Router | None", user_api_key_auth: "UserAPIKeyAuth | None" = None, @@ -306,7 +314,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _route_openai_thinking_to_responses_api_if_needed( completion_kwargs: dict[str, Any], *, - thinking: dict[str, Any] | None, + thinking: Mapping[str, object] | None, ) -> None: """ When users call `litellm.anthropic.messages.*` with a non-Anthropic model and @@ -407,12 +415,12 @@ class LiteLLMMessagesToCompletionTransformationHandler: def _prepare_completion_kwargs( *, max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, - system: str | list[dict[str, Any]] | None = None, + system: _AnthropicSystem = None, temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, @@ -420,7 +428,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, - extra_kwargs: dict[str, Any] | None = None, + extra_kwargs: Mapping[str, object] | None = None, ) -> tuple[dict[str, Any], dict[str, str]]: """Prepare kwargs for litellm.completion/acompletion. @@ -433,7 +441,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: Logging as LiteLLMLoggingObject, ) - request_data: Final = { + request_data: Final[dict[str, object]] = { "model": model, "messages": messages, "max_tokens": max_tokens, @@ -528,7 +536,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, @@ -537,7 +545,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, @@ -551,10 +559,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: requested_router if requested_router is not None else _proxy_router_fallback() ) - proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( - proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None - ) + proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) polyfill_result: Final = await _prepare_context_managed_request( model=model, @@ -618,7 +623,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: _AnthropicMessages, model: str, metadata: dict | None = None, stop_sequences: list[str] | None = None, @@ -627,7 +632,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: temperature: float | None = None, thinking: dict | None = None, tool_choice: dict | None = None, - tools: list[dict] | None = None, + tools: list[dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: dict | None = None, @@ -688,10 +693,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: if context_management is None and not _messages_have_compaction_block(messages): polyfill_result: PolyfillResult | None = None else: - proxy_litellm_metadata: Final = _extract_proxy_litellm_metadata(kwargs) - user_api_key_auth: Final[UserAPIKeyAuth | None] = ( - proxy_litellm_metadata.get("user_api_key_auth") if proxy_litellm_metadata is not None else None - ) + proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs) polyfill_result = run_async_function( _prepare_context_managed_request, model=model, diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index 2d7589a8715..1660f56378f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -348,6 +348,26 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): merged_chunk["context_management"] = ContextManagementResponse(applied_edits=list(self.applied_edits)) return self._augment_message_delta_usage(merged_chunk) + def _handle_choiceless_chunk(self, chunk: "ModelResponseStream") -> bool: + """Consume an OpenAI-compatible chunk that carries no ``choices``. + + ``choices`` is legitimately empty on metadata-only chunks; the final + usage chunk emitted when ``stream_options.include_usage`` is set is the + common case (vLLM and other OpenAI-compatible servers do this). Such a + chunk carries no content-block information, so the caller must not run + the content-block state machine over it. + + Returns True when a merged ``message_delta`` was queued (usage folded + into the held stop-reason chunk); False when the chunk should be + skipped entirely. + """ + if self.holding_stop_reason_chunk is not None and _optional_attr(chunk, "usage") is not None: + self.chunk_queue.append(self._merge_usage_into_held_stop_reason_chunk(chunk)) + self.queued_usage_chunk = True + self.holding_stop_reason_chunk = None + return True + return False + def _ensure_context_management_attached(self, message_delta_chunk: MessageBlockDelta) -> MessageBlockDelta: """Attach ``context_management`` to a ``message_delta`` chunk if ``self.applied_edits`` is non-empty and the chunk does not already @@ -509,6 +529,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if chunk == "None" or chunk is None: raise Exception + if not getattr(chunk, "choices", None): + if self._handle_choiceless_chunk(chunk): + return self.chunk_queue.popleft() + continue + should_start_new_block = self._should_start_new_content_block(chunk) is_opening_first_block = self.sent_content_block_start is False if is_opening_first_block and self._is_blank_delta(chunk): @@ -732,6 +757,11 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if chunk == "None" or chunk is None: raise Exception + if not getattr(chunk, "choices", None): + if self._handle_choiceless_chunk(chunk): + return self.chunk_queue.popleft() + continue + should_start_new_block = self._should_start_new_content_block(chunk) is_opening_first_block = self.sent_content_block_start is False if is_opening_first_block and self._is_blank_delta(chunk): diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 3ef298aa336..c4b5cc628e2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -14,6 +14,7 @@ from typing import Any, Final, cast import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, sanitize_tool_use_ids_in_anthropic_messages, strip_empty_text_blocks_from_anthropic_messages, ) @@ -222,6 +223,7 @@ async def anthropic_messages( # Replay of cross-provider tool history (e.g. kimi -> Anthropic) may carry # ids like ``functions.Bash:0`` that violate Anthropic's id pattern. messages = sanitize_tool_use_ids_in_anthropic_messages(messages) + messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages) from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, @@ -413,6 +415,7 @@ def anthropic_messages_handler( if not kwargs.pop("_litellm_messages_presanitized", False): messages = strip_empty_text_blocks_from_anthropic_messages(messages) messages = sanitize_tool_use_ids_in_anthropic_messages(messages) + messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages) from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, diff --git a/litellm/llms/azure/assistants.py b/litellm/llms/azure/assistants.py index 671e4633af4..f7b419405ac 100644 --- a/litellm/llms/azure/assistants.py +++ b/litellm/llms/azure/assistants.py @@ -1,8 +1,9 @@ from collections.abc import Coroutine, Iterable -from typing import Any, Final, Literal +from typing import Any, Final, Literal, TypedDict import httpx from openai import AsyncAzureOpenAI, AzureOpenAI +from openai.types.shared_params.metadata import Metadata from typing_extensions import overload from ...types.llms.openai import ( @@ -22,6 +23,16 @@ from ...types.llms.openai import ( from .common_utils import BaseAzureLLM +class _RunThreadStreamData(TypedDict): + thread_id: str + assistant_id: str + additional_instructions: str | None + instructions: str | None + metadata: Metadata | None + model: str | None + tools: Iterable[AssistantToolParam] | None + + class AzureAssistantsAPI(BaseAzureLLM): def __init__(self) -> None: super().__init__() @@ -212,9 +223,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj # fmt: off @@ -301,9 +312,9 @@ class AzureAssistantsAPI(BaseAzureLLM): response_obj: OpenAIMessage | None = None if getattr(thread_message, "status", None) is None: thread_message.status = "completed" - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) else: - response_obj = OpenAIMessage(**thread_message.dict()) + response_obj = OpenAIMessage.model_validate(thread_message.dict()) return response_obj async def async_get_messages( @@ -443,7 +454,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = await openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) # fmt: off @@ -539,7 +550,7 @@ class AzureAssistantsAPI(BaseAzureLLM): message_thread: Final = azure_openai_client.beta.threads.create(**data) - return Thread(**message_thread.dict()) + return Thread.model_validate(message_thread.dict()) async def async_get_thread( self, @@ -566,7 +577,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = await openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # fmt: off @@ -642,7 +653,7 @@ class AzureAssistantsAPI(BaseAzureLLM): response: Final = openai_client.beta.threads.retrieve(thread_id=thread_id) - return Thread(**response.dict()) + return Thread.model_validate(response.dict()) # def delete_thread(self): # pass @@ -730,7 +741,8 @@ class AzureAssistantsAPI(BaseAzureLLM): event_handler: AssistantEventHandler | None, litellm_params: dict | None = None, ) -> AssistantStreamManager[AssistantEventHandler]: - data: Final[dict[str, Any]] = { + stream_fn: Final = client.beta.threads.runs.stream + base_data: Final[_RunThreadStreamData] = { "thread_id": thread_id, "assistant_id": assistant_id, "additional_instructions": additional_instructions, @@ -740,8 +752,8 @@ class AzureAssistantsAPI(BaseAzureLLM): "tools": tools, } if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return stream_fn(**base_data, event_handler=event_handler) + return stream_fn(**base_data) # fmt: off diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 5333dec2a20..9e3dec26673 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -12,7 +12,6 @@ from litellm.litellm_core_utils.cloud_storage_security import ( from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret_str from litellm.types.llms.bedrock import ( BedrockCreateBatchRequest, BedrockCreateBatchResponse, @@ -29,7 +28,7 @@ from litellm.types.llms.openai import ( from litellm.types.utils import LiteLLMBatch, LlmProviders from ..base_aws_llm import BaseAWSLLM -from ..common_utils import CommonBatchFilesUtils +from ..common_utils import CommonBatchFilesUtils, resolve_s3_encryption_key_id # Bedrock batch input files are uploaded as # s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see @@ -200,7 +199,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): ) # Add optional KMS encryption key ID if provided - s3_encryption_key_id = litellm_params.get("s3_encryption_key_id") or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + s3_encryption_key_id = resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ) if s3_encryption_key_id: s3_output_config["s3EncryptionKeyId"] = s3_encryption_key_id diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 189bac3256a..d18cb7d8734 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -26,7 +26,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( ) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.secret_managers.main import get_secret +from litellm.secret_managers.main import get_secret, get_secret_str if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues @@ -1304,6 +1304,23 @@ def get_anthropic_beta_from_headers(headers: dict) -> list[str]: return [] +def resolve_s3_encryption_key_id( + litellm_params: Mapping[str, Any], + optional_params: Mapping[str, Any] | None = None, +) -> str | None: + """ + Resolve the SSE-KMS key configured for Bedrock batch/file S3 objects. + + Precedence: `s3_encryption_key_id` in litellm_params, then optional_params + (client-side / request params), then the AWS_S3_ENCRYPTION_KEY_ID env var. + """ + candidates: Final = tuple( + source.get("s3_encryption_key_id") for source in (litellm_params, optional_params) if source is not None + ) + explicit: Final = next((value for value in candidates if isinstance(value, str) and value), None) + return explicit or get_secret_str("AWS_S3_ENCRYPTION_KEY_ID") + + class CommonBatchFilesUtils: """ Common utilities for Bedrock batch and file operations. diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 399f11a94cf..bd3570d50a3 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -2,7 +2,9 @@ import base64 import json import os import time -from collections.abc import Mapping, MutableMapping +from collections.abc import Iterable, Mapping, MutableMapping +from functools import cache +from itertools import chain from types import MappingProxyType from typing import Any, Final from urllib.parse import unquote @@ -10,7 +12,7 @@ from urllib.parse import unquote import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -26,12 +28,16 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) -from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + extract_file_data, + text_completion_prompt_to_messages, +) from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) +from litellm.types.llms.bedrock import BedrockBatchRecordKind from litellm.types.llms.openai import ( AllMessageValues, CreateFileRequest, @@ -41,12 +47,14 @@ from litellm.types.llms.openai import ( OpenAICreateFileRequestOptionalParams, OpenAIFileObject, PathLike, + ResponseInputParam, + ResponsesAPIOptionalRequestParams, ) from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError +from ..common_utils import BedrockError, resolve_s3_encryption_key_id # litellm_params key used to hand the SigV4-signed GET headers from # `transform_file_content_request` to `validate_environment` (the only hook @@ -55,6 +63,26 @@ from ..common_utils import BedrockError S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers" +def _frozen_mapping(items: Iterable[tuple[str, Any]]) -> Mapping[str, Any]: + return MappingProxyType(dict(items)) + + +# JSONL batch records are untyped json, so the `/v1/responses` fields are +# validated into their concrete Responses API types before being handed to the +# Responses-to-Chat bridge. Both adapters drop keys the Responses API doesn't +# define, which is what the bridge would ignore anyway. Built on first use +# rather than at import: `ResponseInputParam` is a deep union and only batch +# files carrying `/v1/responses` records need it. +@cache +def _responses_input_adapter() -> TypeAdapter[str | ResponseInputParam]: + return TypeAdapter(str | ResponseInputParam) + + +@cache +def _responses_request_adapter() -> TypeAdapter[ResponsesAPIOptionalRequestParams]: + return TypeAdapter(ResponsesAPIOptionalRequestParams) + + class _BedrockS3RequestParams(BaseModel): """Typed view of the credential/region params the S3 GetObject path reads.""" @@ -303,41 +331,55 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # example; add others here as they adopt the same schema. CONVERSE_INVOKE_PROVIDERS = ("nova",) - # OpenAI batch URL that signals an embedding request. Per OpenAI Batch API - # spec, every JSONL record carries a `url` field; we use it as the - # authoritative signal to route the line to the embedding code path - # instead of inferring from the presence of `input` vs `messages`. + # OpenAI batch URLs that select which request shape a JSONL line carries. + # Per the OpenAI Batch API spec every record carries a `url`, so we use it + # as the authoritative routing signal instead of inferring from the + # presence of `input` vs `prompt` vs `messages`. OPENAI_EMBEDDINGS_URL = "/v1/embeddings" + OPENAI_TEXT_COMPLETIONS_URL = "/v1/completions" + OPENAI_RESPONSES_URL = "/v1/responses" @staticmethod - def _is_embedding_record(openai_jsonl_record: dict[str, Any]) -> bool: + def _classify_batch_record(openai_jsonl_record: Mapping[str, Any]) -> BedrockBatchRecordKind: """ - Decide whether an OpenAI batch JSONL line is an embedding request. + Decide which OpenAI endpoint shape an OpenAI batch JSONL line carries. - Precedence (strict - any explicit `url` short-circuits): - 1. `url == "/v1/embeddings"` -> embedding. Authoritative per the - OpenAI Batch API spec. - 2. Any other non-empty `url` (e.g. `/v1/chat/completions`) -> NOT - embedding. We trust the caller's explicit signal even if the - body would otherwise suggest embedding; misrouting a chat - record into the embedding transformer would corrupt the - modelInput, while a chat-shaped body sent to the chat path - either succeeds or fails cleanly inside that transformer. - 3. `url` missing/empty -> fall back to body shape. Requires - `input` present AND `messages` absent so a malformed record - carrying both keys routes to the chat path (safer default: - Anthropic transforms ignore unknown top-level keys, whereas - the embedding transformer would silently drop the messages). + Precedence (strict - any recognized `url` short-circuits): + 1. A `url` matching a supported endpoint wins. Authoritative per the + OpenAI Batch API spec, which requires it on every record. + 2. Any other non-empty `url` -> chat. We trust the caller's explicit + signal rather than re-deriving it from the body, and an + unexpectedly-shaped body fails cleanly inside the chat + transformer instead of being silently misrouted. + 3. `url` missing/empty -> fall back to body shape. `messages` wins + over the other keys so a malformed record carrying several of + them keeps its conversation instead of having it dropped, and a + bare `input` stays an embedding for backwards compatibility + (that ambiguity with `/v1/responses` is only resolvable from + `url`). """ - url: Final = openai_jsonl_record.get("url") - if url == BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: - return True - if url: - return False - body: Final = openai_jsonl_record.get("body", {}) - if not isinstance(body, dict): - return False - return "input" in body and "messages" not in body + match openai_jsonl_record.get("url"): + case BedrockFilesConfig.OPENAI_EMBEDDINGS_URL: + return BedrockBatchRecordKind.EMBEDDING + case BedrockFilesConfig.OPENAI_TEXT_COMPLETIONS_URL: + return BedrockBatchRecordKind.TEXT_COMPLETION + case BedrockFilesConfig.OPENAI_RESPONSES_URL: + return BedrockBatchRecordKind.RESPONSES + case None | "": + pass + case _: + return BedrockBatchRecordKind.CHAT + + body: Final = openai_jsonl_record.get("body") + if not isinstance(body, Mapping): + return BedrockBatchRecordKind.CHAT + if "messages" in body: + return BedrockBatchRecordKind.CHAT + if "prompt" in body: + return BedrockBatchRecordKind.TEXT_COMPLETION + if "input" in body: + return BedrockBatchRecordKind.EMBEDDING + return BedrockBatchRecordKind.CHAT # Identifier for the Bedrock Titan v2 InvokeModel body schema as stored # in `model_prices_and_context_window.json`. Centralized so future @@ -544,9 +586,84 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): ) return dict(titan_config._transform_request(input=input_text, inference_params=inference_params)) + @staticmethod + def _transform_text_completion_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Rewrite an OpenAI `/v1/completions` batch body as a Chat Completions body. + + Bedrock batch `modelInput` is the model's InvokeModel/Converse body, and + no Bedrock batch model takes a bare `prompt`, so the wrapping that + `litellm.text_completion` does in real time has to happen here too. + """ + prompt: Final = openai_request_body.get("prompt") + if prompt is None: + raise ValueError( + "Batch record for /v1/completions is missing required `prompt` field: " + f"model={openai_request_body.get('model', '')}" + ) + return _frozen_mapping( + chain( + ((key, value) for key, value in openai_request_body.items() if key != "prompt"), + (("messages", text_completion_prompt_to_messages(prompt)),), + ) + ) + + @staticmethod + def _transform_responses_body_to_chat_body(openai_request_body: Mapping[str, Any]) -> Mapping[str, Any]: + """ + Rewrite an OpenAI `/v1/responses` batch body as a Chat Completions body. + + Delegates to the same Responses-to-Chat bridge the real-time path uses + for providers without a native Responses API (which is every Bedrock + model), so `input`, `instructions`, `max_output_tokens` and the tool + params translate identically in batch and real time. The bridge always + emits a `tools` key; an empty one is dropped rather than shipped as an + empty array inside `modelInput`. + """ + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + responses_input: Final = openai_request_body.get("input") + if responses_input is None: + raise ValueError( + "Batch record for /v1/responses is missing required `input` field: " + f"model={openai_request_body.get('model', '')}" + ) + chat_body: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=openai_request_body.get("model", ""), + input=_responses_input_adapter().validate_python(responses_input), + responses_api_request=_responses_request_adapter().validate_python( + _frozen_mapping( + (key, value) for key, value in openai_request_body.items() if key not in ("model", "input") + ) + ), + metadata=openai_request_body.get("metadata"), + ) + return _frozen_mapping((key, value) for key, value in chat_body.items() if key != "tools" or value) + + @staticmethod + def _transform_batch_body_to_chat_body( + openai_request_body: Mapping[str, Any], + record_kind: BedrockBatchRecordKind, + ) -> Mapping[str, Any]: + """ + Normalize a non-embedding batch body to the Chat Completions shape the + per-provider Bedrock transformations expect. + """ + match record_kind: + case BedrockBatchRecordKind.TEXT_COMPLETION: + return BedrockFilesConfig._transform_text_completion_body_to_chat_body(openai_request_body) + case BedrockBatchRecordKind.RESPONSES: + return BedrockFilesConfig._transform_responses_body_to_chat_body(openai_request_body) + case BedrockBatchRecordKind.CHAT: + return openai_request_body + case BedrockBatchRecordKind.EMBEDDING: + raise ValueError("Embedding batch records do not have a chat-completion equivalent") + def _map_openai_to_bedrock_params( self, - openai_request_body: dict[str, Any], + openai_request_body: Mapping[str, Any], provider: str | None = None, ) -> dict[str, Any]: """ @@ -658,14 +775,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): provider = self.get_bedrock_invoke_provider(model) # Route to the embedding transformer when the OpenAI batch line - # targets /v1/embeddings; otherwise fall back to the existing - # chat-completion path. We branch here (rather than inside + # targets /v1/embeddings; every other endpoint shape is normalized + # to chat completions first. We branch here (rather than inside # `_map_openai_to_bedrock_params`) so the chat helper keeps its # narrow contract and the embedding helper can evolve independently. - if self._is_embedding_record(_openai_jsonl_content): + record_kind = self._classify_batch_record(_openai_jsonl_content) + if record_kind is BedrockBatchRecordKind.EMBEDDING: model_input = self._map_openai_embedding_to_bedrock_params(openai_request_body=openai_body) else: - model_input = self._map_openai_to_bedrock_params(openai_request_body=openai_body, provider=provider) + model_input = self._map_openai_to_bedrock_params( + openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), + provider=provider, + ) # Create Bedrock batch record record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") @@ -733,6 +854,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content=file_content, api_base=api_base, optional_params=optional_params, + s3_encryption_key_id=resolve_s3_encryption_key_id( + litellm_params=litellm_params, + optional_params=optional_params, + ), ) litellm_params["upload_url"] = api_base @@ -750,6 +875,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content: str, api_base: str, optional_params: dict, + s3_encryption_key_id: str | None = None, ) -> tuple[dict, str]: """ Sign S3 PUT request using the same proven logic as S3Logger. @@ -782,12 +908,25 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): content_hash: Final = hashlib.sha256(content.encode("utf-8")).hexdigest() # Prepare headers with required S3 headers (same as s3_v2.py) - request_headers: Final = { - "Content-Type": "application/json", # JSONL files are JSON content - "x-amz-content-sha256": content_hash, # REQUIRED by S3 - "Content-Language": "en", - "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", - } + sse_headers: Final = ( + MappingProxyType( + { + "x-amz-server-side-encryption": "aws:kms", + "x-amz-server-side-encryption-aws-kms-key-id": s3_encryption_key_id, + } + ) + if s3_encryption_key_id + else MappingProxyType({}) + ) + request_headers: Final = MappingProxyType( + { + "Content-Type": "application/json", # JSONL files are JSON content + "x-amz-content-sha256": content_hash, # REQUIRED by S3 + "Content-Language": "en", + "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + **sse_headers, + } + ) # Use requests.Request to prepare the request (same pattern as s3_v2.py) req: Final = requests.Request("PUT", api_base, data=content, headers=request_headers) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index d586156b625..9ada3674d33 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -8,11 +8,12 @@ import sys import threading import time from collections.abc import Callable, Mapping +from http.cookiejar import CookieJar, DefaultCookiePolicy from typing import TYPE_CHECKING, Any, Final, Optional import certifi import httpx -from aiohttp import ClientSession, TCPConnector +from aiohttp import ClientSession, DummyCookieJar, TCPConnector from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport from httpx._types import RequestFiles @@ -144,6 +145,15 @@ def _handler_may_close_client(client_refcount: int, owns_client: bool) -> bool: return owns_client and client_refcount <= _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER +def blocked_cookie_jar() -> CookieJar: + """A jar that stores no response cookie and sends none, for httpx clients. + + LiteLLM's outbound clients are pooled and shared by every caller, so a cookie one + upstream sets would be replayed to every other upstream on a matching domain. + """ + return CookieJar(policy=DefaultCookiePolicy(allowed_domains=())) + + _STREAMING_ERROR_BODY_READ_TIMEOUT_SECONDS: Final = 5.0 _STREAMING_ERROR_BODY_READ_EXECUTOR: Final = concurrent.futures.ThreadPoolExecutor( max_workers=50, @@ -587,6 +597,7 @@ class AsyncHTTPHandler: verify=ssl_config, cert=cert, headers=default_headers, + cookies=blocked_cookie_jar(), follow_redirects=True, ) @@ -1063,6 +1074,7 @@ class AsyncHTTPHandler: def session_factory() -> ClientSession: return ClientSession( connector=TCPConnector(**transport_connector_kwargs), + cookie_jar=DummyCookieJar(), trust_env=trust_env, ) @@ -1132,6 +1144,7 @@ class HTTPHandler: verify=ssl_config, cert=cert, headers=default_headers, + cookies=blocked_cookie_jar(), follow_redirects=True, ) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 29ebd5c2a75..6481b67fad7 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -14,7 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) -from litellm.llms.vertex_ai.common_utils import get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( @@ -98,9 +98,6 @@ class VertexAIBatchPrediction(VertexLLM): data=json.dumps(vertex_batch_request), ) - if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") - _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( response=_json_response @@ -130,8 +127,6 @@ class VertexAIBatchPrediction(VertexLLM): error_body[:1000], ) raise - if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -243,7 +238,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -293,7 +290,9 @@ class VertexAIBatchPrediction(VertexLLM): headers=headers, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -366,7 +365,9 @@ class VertexAIBatchPrediction(VertexLLM): ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response: Final = response.json() vertex_batch_response: Final = ( @@ -391,7 +392,9 @@ class VertexAIBatchPrediction(VertexLLM): params=params, ) if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") + raise VertexAIError( + status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" + ) _json_response: Final = response.json() vertex_batch_response: Final = ( @@ -461,7 +464,7 @@ class VertexAIBatchPrediction(VertexLLM): sync_handler: Final = _get_httpx_client() try: - response: Final = sync_handler.post( + sync_handler.post( url=api_base, headers=headers, data=json.dumps({}), @@ -475,9 +478,6 @@ class VertexAIBatchPrediction(VertexLLM): ) raise - if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") - # HTTPHandler.get() does not accept a timeout parameter retrieve_response: Final = sync_handler.get( url=retrieve_api_base, @@ -489,7 +489,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response: Final = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( @@ -508,7 +511,7 @@ class VertexAIBatchPrediction(VertexLLM): llm_provider=litellm.LlmProviders.VERTEX_AI, ) try: - response: Final = await client.post( + await client.post( url=api_base, headers=headers, data=json.dumps({}), @@ -521,8 +524,6 @@ class VertexAIBatchPrediction(VertexLLM): e.response.text[:1000], ) raise - if response.status_code != 200: - raise Exception(f"Error: {response.status_code} {response.text}") # AsyncHTTPHandler.get() does not accept a timeout parameter retrieve_response: Final = await client.get( @@ -535,7 +536,10 @@ class VertexAIBatchPrediction(VertexLLM): retrieve_response.status_code, retrieve_response.text[:1000], ) - raise Exception(f"Error: {retrieve_response.status_code} {retrieve_response.text}") + raise VertexAIError( + status_code=retrieve_response.status_code, + message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", + ) _json_response: Final = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 66951715488..f284b47292b 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,7 +1,9 @@ from typing import Any, Final +from urllib.parse import unquote from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest @@ -199,16 +201,40 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ - from urllib.parse import unquote - - decoded_uri: Final = unquote(gcs_file_uri) - - model_path: Final = decoded_uri.split("publishers/")[1] - parts: Final = model_path.split("/") - model: Final = f"publishers/{'/'.join(parts[:3])}" + model: Final = cls._parse_model_from_gcs_file(gcs_file_uri) + if model is None: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch creation requires the model to be part of `input_file_id`, but " + f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + "Either upload the input file through LiteLLM (POST /v1/files with " + "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " + "pass a uri of the form " + "gs:////publishers//models//" + ), + ) return model + @classmethod + def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: + """ + Returns the `publishers//models/` path from a gcs uri, or None if the uri + does not contain one. + """ + _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") + if not separator: + return None + + parts: Final = model_path.split("/") + if len(parts) < 3 or parts[1] != "models" or not parts[2]: + return None + + return f"publishers/{'/'.join(parts[:3])}" + @classmethod def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool: """ @@ -216,7 +242,11 @@ class VertexAIBatchTransformation: LiteLLM-managed unified file id) with a `publishers/` model path that `_get_model_from_gcs_file` can parse. """ - return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id + return ( + input_file_id is not None + and input_file_id.startswith("gs://") + and cls._parse_model_from_gcs_file(input_file_id) is not None + ) @classmethod def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 76618e0f742..e285feb77ee 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -109,7 +109,7 @@ if MCP_AVAILABLE: ############ MCP Server REST API Routes ################# async def _safe_fire_mcp_tool_call_logging( logging_obj: Any | None, - result: Any, + result: "CallToolResult", start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 1c6ad84ddb4..49a1f1314f0 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -13,9 +13,9 @@ import time import traceback import types import uuid -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from fastapi import FastAPI, HTTPException @@ -145,7 +145,7 @@ try: ) # Robust auth lookup keyed by session_object. - _session_obj_auth_storage: "weakref.WeakKeyDictionary[Any, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() + _session_obj_auth_storage: "weakref.WeakKeyDictionary[object, MCPAuthenticatedUser]" = weakref.WeakKeyDictionary() except ImportError as e: verbose_logger.debug("MCP module not found: %s", e) MCP_AVAILABLE = False @@ -493,14 +493,14 @@ if MCP_AVAILABLE: def _gateway_create_initialization_options( self, notification_options: NotificationOptions | None = None, - experimental_capabilities: dict[str, dict[str, Any]] | None = None, + experimental_capabilities: dict[str, dict[str, object]] | None = None, ) -> InitializationOptions: opts: Final = Server.create_initialization_options( self, notification_options=notification_options, experimental_capabilities=experimental_capabilities or {}, ) - updates: Final[dict[str, Any]] = {} + updates: Final[dict[str, str]] = {} merged: Final = _mcp_gateway_initialize_instructions.get() if merged is not None: updates["instructions"] = merged @@ -549,6 +549,17 @@ if MCP_AVAILABLE: _stateful_session_locks: Final[dict[str, asyncio.Lock]] = {} _stateful_session_active_request_counts: Final[dict[str, int]] = {} + class _TerminableTransport(Protocol): + async def terminate(self) -> None: ... + + class _TransportRegistry(Protocol): + def __contains__(self, session_id: object, /) -> bool: ... + + def pop(self, session_id: str, default: None, /) -> "_TerminableTransport | None": ... + + def _stateful_server_instances() -> _TransportRegistry: + return getattr(session_manager_stateful, "_server_instances", {}) + def _remove_stateful_session_tracking(session_id: str) -> None: _stateful_session_auth_contexts.pop(session_id, None) _stateful_session_auth_context_last_seen.pop(session_id, None) @@ -578,8 +589,8 @@ if MCP_AVAILABLE: ) -> None: """Terminate expired stateful sessions and drop their auth contexts.""" now = time.monotonic() if now is None else now - server_instances: Final = getattr(session_manager_stateful, "_server_instances", {}) - expired_session_ids: Final = [] + server_instances: Final = _stateful_server_instances() + expired_session_ids: Final[list[str]] = [] for session_id, last_seen in _stateful_session_auth_context_last_seen.items(): if _stateful_session_active_request_counts.get(session_id, 0) > 0: continue @@ -619,7 +630,7 @@ if MCP_AVAILABLE: session may proceed, or ``False`` when the caller is already at the cap with every session in flight (the new ``initialize`` should be rejected). """ - server_instances: Final = getattr(session_manager_stateful, "_server_instances", {}) + server_instances: Final = _stateful_server_instances() def _owned_live_session_ids() -> list[str]: return [ @@ -778,7 +789,7 @@ if MCP_AVAILABLE: get_virtual_tool_definitions, ) - return [Tool(**d) for d in get_virtual_tool_definitions()] + return [Tool.model_validate(d) for d in get_virtual_tool_definitions()] # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") @@ -847,7 +858,7 @@ if MCP_AVAILABLE: async def _build_virtual_call_logging_obj( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth, ) -> LiteLLMLoggingObj | None: """Run the pre-call pipeline (guardrails + logging setup) for a virtual @@ -885,7 +896,7 @@ if MCP_AVAILABLE: async def _dispatch_virtual_mcp_tool( name: str, - arguments: dict[str, Any] | None, + arguments: dict[str, object] | None, user_api_key_auth: UserAPIKeyAuth | None, client_ip: str | None, mcp_servers: list[str] | None = None, @@ -957,7 +968,7 @@ if MCP_AVAILABLE: ) @server.call_tool() - async def mcp_server_tool_call(name: str, arguments: dict[str, Any] | None) -> CallToolResult: + async def mcp_server_tool_call(name: str, arguments: dict[str, object] | None) -> CallToolResult: """ Call a specific tool with the provided arguments Args: @@ -1621,7 +1632,7 @@ if MCP_AVAILABLE: async def _get_user_oauth_extra_headers_from_db( server: MCPServer, user_api_key_auth: UserAPIKeyAuth | None, - prefetched_creds: dict[str, dict[str, Any]] | None = None, + prefetched_creds: 'Mapping[str, "OAuthCredentialPayload"] | None' = None, ) -> dict[str, str] | None: """Stored OAuth2 token for (user, server) as an ``Authorization: Bearer`` header, or None. @@ -1646,7 +1657,7 @@ if MCP_AVAILABLE: Returns a dict keyed by server_id to avoid N+1 queries in asyncio.gather loops. """ - user_id: Final = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) if user_api_key_auth else None if not user_id: return {} try: @@ -1871,7 +1882,7 @@ if MCP_AVAILABLE: list_tools_start_time: Final = datetime.now() litellm_logging_obj: LiteLLMLoggingObj | None = None - list_tools_request_data: dict[str, Any] = {} + list_tools_request_data: dict[str, object] = {} if log_list_tools_to_spendlogs: # This is intentionally minimal: only async_success_handler / post_call_failure_hook @@ -1879,7 +1890,7 @@ if MCP_AVAILABLE: list_tools_call_id: Final = str(uuid.uuid4()) # Derive trace_id from raw_headers when not explicitly passed (same as A2A / MCP call_tool) effective_litellm_trace_id: Final = litellm_trace_id or get_chain_id_from_headers(raw_headers) - spend_logs_metadata: Final[dict[str, Any]] = { + spend_logs_metadata: Final[dict[str, object]] = { "mcp_operation": "list_tools", } if isinstance(list_tools_log_source, str): @@ -2615,7 +2626,7 @@ if MCP_AVAILABLE: async def execute_mcp_tool( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], allowed_mcp_servers: list[MCPServer], start_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -2882,7 +2893,7 @@ if MCP_AVAILABLE: _request_auth_header.reset(_auth_token) _request_extra_headers.reset(_extra_token) _request_resolved_auth_headers.reset(_resolved_token) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) # Try managed MCP server tool (the name is bare; the prefix boundary was # already resolved above against this server's registered prefixes) @@ -2956,7 +2967,7 @@ if MCP_AVAILABLE: arguments = hook_result["arguments"] # pyright: ignore[reportAny] # hook returns untyped args local_content = await _handle_local_mcp_tool(original_tool_name, arguments) - response = CallToolResult(content=cast(Any, local_content), isError=False) + response = CallToolResult(content=local_content, isError=False) return await _run_post_mcp_call_guardrails( result=response, @@ -3003,7 +3014,7 @@ if MCP_AVAILABLE: async def _fire_mcp_tool_call_logging( logging_obj: LiteLLMLoggingObj, - result: Any, + result: CallToolResult, start_time: datetime, end_time: datetime, user_api_key_auth: UserAPIKeyAuth | None = None, @@ -3070,7 +3081,7 @@ if MCP_AVAILABLE: @client async def call_mcp_tool( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3161,7 +3172,7 @@ if MCP_AVAILABLE: async def mcp_get_prompt( name: str, - arguments: dict[str, Any] | None = None, + arguments: dict[str, object] | None = None, user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_servers: list[str] | None = None, @@ -3262,7 +3273,7 @@ if MCP_AVAILABLE: def _get_standard_logging_mcp_tool_call( name: str, - arguments: dict[str, Any], + arguments: dict[str, object], server_name: str | None, session_id: str | None = None, ) -> StandardLoggingMCPToolCall: @@ -3291,13 +3302,13 @@ if MCP_AVAILABLE: async def _handle_managed_mcp_tool( server_name: str, name: str, - arguments: dict[str, Any], + arguments: dict[str, object], user_api_key_auth: UserAPIKeyAuth | None = None, mcp_auth_header: str | None = None, mcp_server_auth_headers: dict[str, dict[str, str]] | None = None, oauth2_headers: dict[str, str] | None = None, raw_headers: dict[str, str] | None = None, - litellm_logging_obj: Any | None = None, + litellm_logging_obj: LiteLLMLoggingObj | None = None, host_progress_callback: Callable | None = None, ) -> CallToolResult: """Handle tool execution for managed server tools""" @@ -3320,7 +3331,7 @@ if MCP_AVAILABLE: return call_tool_result async def _handle_local_mcp_tool( - name: str, arguments: dict[str, Any] + name: str, arguments: dict[str, object] ) -> list[TextContent | ImageContent | EmbeddedResource]: """ Handle tool execution for local registry tools @@ -3426,7 +3437,8 @@ if MCP_AVAILABLE: Extract mcp-session-id from ASGI scope headers. Returns None if not present. """ - for header_name, header_value in scope.get("headers", []): + scope_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) + for header_name, header_value in scope_headers: name = header_name if isinstance(header_name, bytes) else header_name.encode() if name.lower() == b"mcp-session-id": return header_value.decode() if isinstance(header_value, bytes) else str(header_value) @@ -3528,7 +3540,7 @@ if MCP_AVAILABLE: if message.get("type") != "http.request": break - body = message.get("body", b"") or b"" + body: bytes = message.get("body", b"") or b"" if body: # Only retain up to the remaining peek budget for sniffing. # The full ``message`` is already in memory (delivered by @@ -3571,9 +3583,9 @@ if MCP_AVAILABLE: Fixes https://github.com/BerriAI/litellm/issues/20992 """ _mcp_session_header: Final = b"mcp-session-id" - _headers: Final = scope.get("headers", []) + _headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = scope.get("headers", []) - def _normalize_header_name(header_name: Any) -> bytes | None: + def _normalize_header_name(header_name: object) -> bytes | None: if isinstance(header_name, bytes): return header_name.lower() if isinstance(header_name, str): @@ -3902,7 +3914,8 @@ if MCP_AVAILABLE: def _get_authorization_header_from_scope(scope: Scope) -> str | None: """First ``Authorization`` header value in the ASGI scope, or None.""" - for key, value in scope.get("headers", []): + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + for key, value in scope_headers: if key.lower() == b"authorization": return value.decode("latin-1") return None @@ -3921,7 +3934,8 @@ if MCP_AVAILABLE: ``MCPRequestHandler.process_mcp_request``), and forwarding it upstream would leak the proxy key to a third-party MCP server. """ - has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope.get("headers", [])) + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + has_litellm_key_header: Final = any(key.lower() == b"x-litellm-api-key" for key, _ in scope_headers) if not has_litellm_key_header: return None return _get_authorization_header_from_scope(scope) @@ -4115,7 +4129,7 @@ if MCP_AVAILABLE: async def handle_streamable_http_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through StreamableHTTP.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, @@ -4135,7 +4149,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar (set by # /toolset/{name}/mcp and /{name}/mcp route handlers in proxy_server.py). @@ -4436,7 +4451,7 @@ if MCP_AVAILABLE: async def handle_sse_mcp(scope: Scope, receive: Receive, send: Send) -> None: """Handle MCP requests through SSE.""" try: - path: Final = scope.get("path", "") + path: Final[str] = scope.get("path", "") ( user_api_key_auth, mcp_auth_header, @@ -4456,7 +4471,8 @@ if MCP_AVAILABLE: ) # Strip any client-supplied x-mcp-toolset-id to prevent forgery. - scope["headers"] = [(k, v) for k, v in scope.get("headers", []) if k.lower() != b"x-mcp-toolset-id"] + scope_headers: Final[Sequence[tuple[bytes, bytes]]] = scope.get("headers", []) + scope["headers"] = [(k, v) for k, v in scope_headers if k.lower() != b"x-mcp-toolset-id"] # Apply toolset scope if set server-side via ContextVar so the # downstream probe list matches the fully-authorized server set @@ -4680,7 +4696,8 @@ if MCP_AVAILABLE: ) -> Send: async def wrapped_send(message: Message) -> None: if message.get("type") == "http.response.start": - for key, value in message.get("headers", []): + response_headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = message.get("headers", []) + for key, value in response_headers: header_name = key if isinstance(key, bytes) else str(key).encode() if header_name.lower() == b"mcp-session-id": session_id = value.decode() if isinstance(value, bytes) else str(value) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5b1134650f2..1fc05ac4653 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2488,6 +2488,16 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "is active as a reminder that hard enforcement is relaxed." ), ) + apply_user_budget_to_team_keys: bool | None = Field( + None, + description=( + "If True, a user's personal max_budget is enforced on every request they " + "make, including requests made with a team-scoped key. Defaults to False, " + "where a team-scoped key is governed only by the team and team-member " + "budgets and the key owner's personal max_budget does not apply " + "(see GitHub issue #12905)." + ), + ) user_url_validation: bool | None = Field( None, description=( @@ -3884,12 +3894,19 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse): ########################################## +class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): + access_group_id: str + access_group_name: str + models: tuple[str, ...] + + class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): team_member_budget_table: LiteLLM_BudgetTableFull | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None access_group_agent_ids: list[str] | None = None + access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 476bd725c73..742fdf35b1e 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -1,8 +1,10 @@ +import asyncio import hashlib import json from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, Protocol, TypedDict +from types import MappingProxyType +from typing import Any, Final, NamedTuple, Protocol, TypedDict import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -10,7 +12,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) from litellm.proxy.utils import PrismaClient -from litellm.repositories.table_repositories import AgentsRepository +from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -86,10 +88,32 @@ def agents_table(prisma_client: PrismaClient) -> AgentTableClient: return table +class ObjectPermissionGrantRecord(Protocol): + object_permission_id: str + agents: list[str] | None + + +class ObjectPermissionTableClient(Protocol): + async def find_many(self, where: Mapping[str, object]) -> Sequence[ObjectPermissionGrantRecord]: ... + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient: + table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table + return table + + +class GrantMigrationResult(NamedTuple): + rewritten: int + missed: int + + class AgentRegistry: def __init__(self): self.agent_list: list[AgentResponse] = [] self.config_agents: tuple[AgentConfig, ...] = () + self.config_agent_legacy_ids: Mapping[str, str] = MappingProxyType({}) def reset_agent_list(self): self.agent_list = [] @@ -100,23 +124,33 @@ class AgentRegistry: def deregister_agent(self, agent_name: str): self.agent_list = [agent for agent in self.agent_list if agent.agent_name != agent_name] - def get_agent_list(self, agent_names: Sequence[str] | None = None): + def get_agent_list(self, agent_names: Sequence[str] | None = None) -> tuple[AgentResponse, ...]: if agent_names is not None: - return [agent for agent in self.agent_list if agent.agent_name in agent_names] - return self.agent_list + return tuple(agent for agent in self.agent_list if agent.agent_name in agent_names) + return tuple(self.agent_list) - def get_public_agent_list(self) -> list[AgentResponse]: - public_agent_list: Final[list[AgentResponse]] = [] - if litellm.public_agent_groups is None: - return public_agent_list - for agent in self.agent_list: - if agent.agent_id in litellm.public_agent_groups: - public_agent_list.append(agent) - return public_agent_list + def get_public_agent_list(self) -> tuple[AgentResponse, ...]: + public_agent_groups: Final = litellm.public_agent_groups + if public_agent_groups is None: + return () + return tuple( + agent for agent in self.agent_list if not self.ids_for_agent(agent.agent_id).isdisjoint(public_agent_groups) + ) def _create_agent_id(self, agent_config: AgentConfig) -> str: + return hashlib.sha256(agent_config["agent_name"].encode()).hexdigest() + + def _create_legacy_agent_id(self, agent_config: AgentConfig) -> str: return hashlib.sha256(json.dumps(agent_config, sort_keys=True).encode()).hexdigest() + def ids_for_agent(self, agent_id: str) -> frozenset[str]: + return frozenset( + {agent_id, *(legacy for legacy, stable in self.config_agent_legacy_ids.items() if stable == agent_id)} + ) + + def stable_agent_id(self, agent_id: str) -> str: + return self.config_agent_legacy_ids.get(agent_id, agent_id) + def load_agents_from_config(self, agent_config: Sequence[AgentConfig] | None = None): """ Register the agents declared in config.yaml and remember them for later rebuilds. @@ -131,12 +165,20 @@ class AgentRegistry: if agent_config is None: return - self.config_agents = tuple(agent_config) - for agent_config_item in agent_config: if not isinstance(agent_config_item, dict): raise ValueError("agent_config must be a list of dictionaries") + self.config_agents = tuple(agent_config) + self.config_agent_legacy_ids = MappingProxyType( + { + self._create_legacy_agent_id(agent_config_item): self._create_agent_id(agent_config_item) + for agent_config_item in agent_config + if agent_config_item.get("agent_name") and agent_config_item.get("agent_card_params") + } + ) + + for agent_config_item in agent_config: agent_name = agent_config_item.get("agent_name") agent_card_params = agent_config_item.get("agent_card_params") if not all([agent_name, agent_card_params]): @@ -180,6 +222,45 @@ class AgentRegistry: self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents) return self.agent_list + async def migrate_legacy_grant_ids(self, table: ObjectPermissionTableClient) -> GrantMigrationResult: + """ + Rewrite object_permission.agents rows holding a legacy full-entry hash to the + stable name-derived id. + + Only the running proxy can do this: the legacy hash is computed from the + resolved config entry (secrets included), so no SQL migration can know it. + Persisting the stable id here is what keeps a grant alive across a later + secret rotation, which re-mints the legacy hash and would otherwise orphan + the stored value. Idempotent; runs of it after the first find no rows. + + Each write is a compare-and-swap against the agents array read above, so a + grant edited concurrently is left untouched; the runtime alias keeps covering + it and the next boot retries the rewrite. + """ + legacy_ids: Final = tuple(legacy for legacy, stable in self.config_agent_legacy_ids.items() if legacy != stable) + if not legacy_ids: + return GrantMigrationResult(rewritten=0, missed=0) + rows: Final = await table.find_many(where={"agents": {"has_some": legacy_ids}}) + updates: Final = tuple( + ( + row.object_permission_id, + tuple(row.agents or ()), + tuple(dict.fromkeys(self.stable_agent_id(agent_id) for agent_id in row.agents or ())), + ) + for row in rows + ) + counts: Final = await asyncio.gather( + *( + table.update_many( + where={"object_permission_id": object_permission_id, "agents": {"equals": snapshot_agents}}, + data={"agents": translated_agents}, + ) + for object_permission_id, snapshot_agents, translated_agents in updates + ) + ) + rewritten: Final = sum(counts) + return GrantMigrationResult(rewritten=rewritten, missed=len(updates) - rewritten) + ########################################################### ########### DB management helpers for agents ########### ############################################################ @@ -492,6 +573,14 @@ class AgentRegistry: if agent.agent_id == agent_id: return agent + translated_id: Final = self.config_agent_legacy_ids.get(agent_id) + if translated_id is None: + return None + + for agent in self.agent_list: + if agent.agent_id == translated_id: + return agent + return None except Exception as e: raise Exception(f"Error getting agent from DB: {e}") diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index af7e731b385..81586b4eef3 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -5,7 +5,8 @@ Handles agent permission checking for keys and teams using object_permission_id. Follows the same pattern as MCP permission handling. """ -from typing import Final +from dataclasses import dataclass +from typing import Final, TypeAlias from litellm._logging import verbose_logger from litellm.proxy._types import ( @@ -17,6 +18,27 @@ from litellm.proxy._types import ( from litellm.repositories.table_repositories import AgentsRepository +@dataclass(frozen=True, slots=True) +class UnrestrictedAgentAccess: + """No agent grant exists on the key or its team, so every agent is reachable.""" + + +@dataclass(frozen=True, slots=True) +class RestrictedAgentAccess: + """Only ``agent_ids`` are reachable. An empty set denies every agent.""" + + agent_ids: frozenset[str] + + +AgentAccess: TypeAlias = UnrestrictedAgentAccess | RestrictedAgentAccess + + +def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + return frozenset(global_agent_registry.stable_agent_id(agent_id) for agent_id in agent_ids) + + class AgentRequestHandler: """ Class to handle agent permission checking, including: @@ -32,37 +54,32 @@ class AgentRequestHandler: """ @staticmethod - async def get_allowed_agents( + async def resolve_agent_access( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ - Get list of allowed agent IDs for the given user/key based on permissions. + Resolve the agents the given user/key may reach. - Returns: - List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all). + ``UnrestrictedAgentAccess`` is only returned when neither the key nor its team + carries any grant. Grants that intersect to nothing stay restricted, so + narrowing a caller can never widen what it reaches. """ try: - allowed_agents: list[str] = [] - allowed_agents_for_key: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) - allowed_agents_for_team: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) + key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) + team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) - # If team has agent restrictions, handle inheritance and intersection logic - if len(allowed_agents_for_team) > 0: - if len(allowed_agents_for_key) > 0: - # Key has its own agent permissions - use intersection with team permissions - for agent_id in allowed_agents_for_key: - if agent_id in allowed_agents_for_team: - allowed_agents.append(agent_id) - else: - # Key has no agent permissions - inherit from team - allowed_agents = allowed_agents_for_team - else: - allowed_agents = allowed_agents_for_key - - return list(set(allowed_agents)) + match (key_access, team_access): + case (UnrestrictedAgentAccess(), UnrestrictedAgentAccess()): + return UnrestrictedAgentAccess() + case (UnrestrictedAgentAccess(), RestrictedAgentAccess(team_ids)): + return RestrictedAgentAccess(_to_stable_ids(team_ids)) + case (RestrictedAgentAccess(key_ids), UnrestrictedAgentAccess()): + return RestrictedAgentAccess(_to_stable_ids(key_ids)) + case (RestrictedAgentAccess(key_ids), RestrictedAgentAccess(team_ids)): + return RestrictedAgentAccess(_to_stable_ids(key_ids) & _to_stable_ids(team_ids)) except Exception as e: verbose_logger.warning("Failed to get allowed agents: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod async def is_agent_allowed( @@ -79,13 +96,14 @@ class AgentRequestHandler: Returns: bool: True if agent is allowed, False otherwise """ - allowed_agents: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth) + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - # Empty list means no restrictions - allow all - if len(allowed_agents) == 0: - return True - - return agent_id in allowed_agents + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth): + case UnrestrictedAgentAccess(): + return True + case RestrictedAgentAccess(allowed_agent_ids): + stable_id: Final = global_agent_registry.stable_agent_id(agent_id) + return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agent_ids) @staticmethod def _get_key_object_permission( @@ -139,55 +157,58 @@ class AgentRequestHandler: @staticmethod async def _get_allowed_agents_for_key( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ Get allowed agents for a key. 1. First checks native key-level agent permissions (object_permission) 2. Also includes agents from key's access_group_ids (unified access groups) + A key that declares agents or access groups is restricted even when those + declarations resolve to nothing, so an emptied or deleted access group denies + rather than opening the key up. Lookup failures still propagate to the caller, + which keeps them fail-open. + Note: object_permission is already loaded by get_key_object() in main auth flow. """ if user_api_key_auth is None: - return [] + return UnrestrictedAgentAccess() try: - all_agents: list[str] = [] - # 1. Get agents from object_permission (native permissions) key_object_permission: Final = AgentRequestHandler._get_key_object_permission(user_api_key_auth) - if key_object_permission is not None: - # Get direct agents - direct_agents: Final = key_object_permission.agents or [] - - # Get agents from access groups - access_group_agents: Final = await AgentRequestHandler._get_agents_from_access_groups( - key_object_permission.agent_access_groups or [] - ) - - all_agents = direct_agents + access_group_agents - + direct_agents: Final = tuple( + key_object_permission.agents or () if key_object_permission is not None else () + ) + declared_access_groups: Final = tuple( + key_object_permission.agent_access_groups or () if key_object_permission is not None else () + ) # 2. Fallback: get agent IDs from key's access_group_ids (unified access groups) - key_access_group_ids: Final = user_api_key_auth.access_group_ids or [] - if key_access_group_ids: - from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - ) + key_access_group_ids: Final = tuple(user_api_key_auth.access_group_ids or ()) - unified_agents: Final = await _get_agent_ids_from_access_groups( - access_group_ids=key_access_group_ids, - ) - all_agents.extend(unified_agents) + if not direct_agents and not declared_access_groups and not key_access_group_ids: + return UnrestrictedAgentAccess() - return list(set(all_agents)) + access_group_agents: Final = ( + tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) + if declared_access_groups + else () + ) + unified_agents: Final = ( + tuple(await AgentRequestHandler._get_unified_access_group_agents(list(key_access_group_ids))) + if key_access_group_ids + else () + ) + + return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) except Exception as e: verbose_logger.warning("Failed to get allowed agents for key: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod async def _get_allowed_agents_for_team( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ Get allowed agents for a team. @@ -195,12 +216,13 @@ class AgentRequestHandler: 2. Also includes agents from team's access_group_ids (unified access groups) Fetches the team object once and reuses it for both permission sources. + Declared-but-empty grants stay restricted; see `_get_allowed_agents_for_key`. """ if user_api_key_auth is None: - return [] + return UnrestrictedAgentAccess() if user_api_key_auth.team_id is None: - return [] + return UnrestrictedAgentAccess() try: from litellm.proxy.auth.auth_checks import get_team_object @@ -211,7 +233,7 @@ class AgentRequestHandler: ) if not prisma_client: - return [] + return UnrestrictedAgentAccess() # Fetch the team object once for both permission sources team_obj: Final = await get_team_object( @@ -223,42 +245,38 @@ class AgentRequestHandler: ) if team_obj is None: - return [] - - all_agents: list[str] = [] + return UnrestrictedAgentAccess() # 1. Get agents from object_permission (native permissions) object_permissions: Final = team_obj.object_permission - if object_permissions is not None: - # Get direct agents - direct_agents: Final = object_permissions.agents or [] - - # Get agents from access groups - access_group_agents: Final = await AgentRequestHandler._get_agents_from_access_groups( - object_permissions.agent_access_groups or [] - ) - - all_agents = direct_agents + access_group_agents - + direct_agents: Final = tuple(object_permissions.agents or () if object_permissions is not None else ()) + declared_access_groups: Final = tuple( + object_permissions.agent_access_groups or () if object_permissions is not None else () + ) # 2. Also include agents from team's access_group_ids (unified access groups) - team_access_group_ids: Final = team_obj.access_group_ids or [] - if team_access_group_ids: - from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - ) + team_access_group_ids: Final = tuple(team_obj.access_group_ids or ()) - unified_agents: Final = await _get_agent_ids_from_access_groups( - access_group_ids=team_access_group_ids, - ) - all_agents.extend(unified_agents) + if not direct_agents and not declared_access_groups and not team_access_group_ids: + return UnrestrictedAgentAccess() - return list(set(all_agents)) + access_group_agents: Final = ( + tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) + if declared_access_groups + else () + ) + unified_agents: Final = ( + tuple(await AgentRequestHandler._get_unified_access_group_agents(list(team_access_group_ids))) + if team_access_group_ids + else () + ) + + return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) except Exception as e: # litellm-dashboard is the default UI team and will never have agents; # skip noisy warnings for it. if user_api_key_auth.team_id != UI_TEAM_ID: verbose_logger.warning("Failed to get allowed agents for team: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod def _get_config_agent_ids_for_access_groups(config_agents: list, access_groups: list[str]) -> set[str]: @@ -277,18 +295,26 @@ class AgentRequestHandler: async def _get_db_agent_ids_for_access_groups(prisma_client, access_groups: list[str]) -> set[str]: """ Helper to get agent_ids from DB agents that match any of the given access groups. + + Query failures propagate so the caller can tell "this group is empty" (deny) + apart from "the lookup failed" (fail-open). """ - agent_ids: Final[set[str]] = set() - if access_groups and prisma_client is not None: - try: - agents: Final = await AgentsRepository(prisma_client).table.find_many( - where={"agent_access_groups": {"hasSome": access_groups}} - ) - for agent in agents: - agent_ids.add(agent.agent_id) - except Exception as e: - verbose_logger.debug("Error getting agents from access groups: %s", e) - return agent_ids + if not access_groups or prisma_client is None: + return set() + + agents: Final = await AgentsRepository(prisma_client).table.find_many( + where={"agent_access_groups": {"hasSome": access_groups}} + ) + return {agent.agent_id for agent in agents} + + @staticmethod + async def _get_unified_access_group_agents(access_group_ids: list[str]) -> list[str]: + """ + Resolve unified access group ids to agent IDs. + """ + from litellm.proxy.auth.auth_checks import _get_agent_ids_from_access_groups + + return await _get_agent_ids_from_access_groups(access_group_ids=access_group_ids) @staticmethod async def _get_agents_from_access_groups( @@ -300,20 +326,17 @@ class AgentRequestHandler: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.proxy_server import prisma_client - try: - # Use the helper for config-loaded agents - agent_ids: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups( - global_agent_registry.agent_list, access_groups - ) + # Use the helper for config-loaded agents + config_agent_ids: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups( + global_agent_registry.agent_list, access_groups + ) - # Use the helper for DB agents - db_agent_ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups(prisma_client, access_groups) - agent_ids.update(db_agent_ids) + # Use the helper for DB agents + db_agent_ids: Final = await AgentRequestHandler._get_db_agent_ids_for_access_groups( + prisma_client, access_groups + ) - return list(agent_ids) - except Exception as e: - verbose_logger.warning("Failed to get agents from access groups: %s", e) - return [] + return list(config_agent_ids | db_agent_ids) @staticmethod async def get_agent_access_groups( diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 1f9c6e1cc05..d348bc01153 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -101,7 +101,11 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) foreign key. Mirrors how spend is joined into the agent response so the UI never has to cross-reference a full key dump client-side. Only non-secret fields are exposed (alias, masked key_name, hashed token).""" - agent_ids: Final = [agent.agent_id for agent in agents] + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + agent_ids: Final = tuple( + alias_id for agent in agents for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + ) if not agent_ids: return key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many( @@ -117,7 +121,12 @@ async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) ) ) for agent in agents: - agent.keys = keys_by_agent.get(agent.agent_id) + matched_keys = [ + key_summary + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + for key_summary in keys_by_agent.get(alias_id) or () + ] + agent.keys = matched_keys or None def _redact_sensitive_agent_fields( @@ -239,10 +248,12 @@ async def get_agents( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) try: - returned_agents: list[AgentResponse] = [] + returned_agents: Sequence[AgentResponse] = () # Admin users get all agents if ( @@ -252,37 +263,45 @@ async def get_agents( returned_agents = global_agent_registry.get_agent_list() else: # Get allowed agents from object_permission (key/team level) - allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) + agent_access: Final = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict) + all_agents: Final = global_agent_registry.get_agent_list() - # If no restrictions (empty list), return all agents - if len(allowed_agent_ids) == 0: - returned_agents = global_agent_registry.get_agent_list() - else: - # Filter agents by allowed IDs - all_agents: Final = global_agent_registry.get_agent_list() - returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] + match agent_access: + case UnrestrictedAgentAccess(): + returned_agents = all_agents + case RestrictedAgentAccess(allowed_agent_ids): + returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] # Fetch current spend from DB for all returned agents from litellm.proxy.proxy_server import prisma_client if prisma_client is not None: - agent_ids: Final = [agent.agent_id for agent in returned_agents] + agent_ids: Final = tuple( + alias_id + for agent in returned_agents + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + ) if agent_ids: db_agents: Final = await agents_table(prisma_client).find_many( where={"agent_id": {"in": agent_ids}}, ) spend_map: Final = {a.agent_id: a.spend for a in db_agents} for agent in returned_agents: - if agent.agent_id in spend_map: - agent.spend = spend_map[agent.agent_id] + matched_spends = tuple( + spend_map[alias_id] + for alias_id in global_agent_registry.ids_for_agent(agent.agent_id) + if alias_id in spend_map + ) + if matched_spends: + agent.spend = sum(matched_spends) await _attach_keys_to_agents(returned_agents, prisma_client) # add is_public field to each agent - we do it this way, to allow setting config agents as public for agent in returned_agents: if agent.litellm_params is None: agent.litellm_params = {} - agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and ( - agent.agent_id in litellm.public_agent_groups + agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and not ( + global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups) ) # Redact sensitive fields for non-admin users @@ -863,7 +882,7 @@ async def make_agent_public( if litellm.public_agent_groups is None: litellm.public_agent_groups = [] # handle duplicates - if agent.agent_id in litellm.public_agent_groups: + if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups): raise HTTPException( status_code=400, detail=f"Agent with name {agent.agent_name} already in public agent groups", @@ -1043,27 +1062,29 @@ async def get_agent_daily_activity( # intersect their explicit `agent_ids` filter with the same allowlist. from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view where_condition: Final[dict[str, object]] = {} if not _user_has_admin_view(user_api_key_dict): - permitted_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - # `get_allowed_agents` returns an empty list when the caller's key - # and team carry no agent restrictions. For activity scoping that's - # not "see everything" — fall back to the agents the caller - # created so they cannot enumerate other tenants' agents. + permitted_agent_ids: list[str] = [] + # An unrestricted caller is not "see everything" for activity scoping. Fall + # back to the agents the caller created so they cannot enumerate other + # tenants' agents. # Guard against `user_id is None`: a literal None in Prisma # `where={"created_by": None}` resolves to ``created_by IS NULL`` # and would expose every ownerless agent's rows. - if not permitted_agent_ids: - if user_api_key_dict.user_id is None: - permitted_agent_ids = [] - else: - owned_records: Final = await agents_table(prisma_client).find_many( - where={"created_by": user_api_key_dict.user_id} - ) - permitted_agent_ids = [a.agent_id for a in owned_records] + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + permitted_agent_ids = list(allowed_agent_ids) + case UnrestrictedAgentAccess(): + if user_api_key_dict.user_id is not None: + owned_records: Final = await agents_table(prisma_client).find_many( + where={"created_by": user_api_key_dict.user_id} + ) + permitted_agent_ids = [a.agent_id for a in owned_records] if agent_ids_list: permitted_agent_id_set: Final = set(permitted_agent_ids) diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 862b66d3860..4a88644bae5 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -4,8 +4,6 @@ Helper functions for appending A2A agents to model lists. Used by proxy model endpoints to make agents appear in UI alongside models. """ -from typing import Final - from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -27,20 +25,23 @@ async def append_agents_to_model_group( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, ) - allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - - for agent_id in allowed_agent_ids: - agent = global_agent_registry.get_agent_by_id(agent_id) - if agent is not None: - model_groups.append( - ModelGroupInfoProxy( - model_group=f"a2a/{agent.agent_name}", - mode="chat", - providers=["a2a"], - ) - ) + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + for agent_id in allowed_agent_ids: + agent = global_agent_registry.get_agent_by_id(agent_id) + if agent is not None: + model_groups.append( + ModelGroupInfoProxy( + model_group=f"a2a/{agent.agent_name}", + mode="chat", + providers=["a2a"], + ) + ) + case _: + pass except Exception as e: verbose_proxy_logger.debug("Error appending agents to model_group/info: %s", e) @@ -61,30 +62,33 @@ async def append_agents_to_model_info( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, ) - allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - - for agent_id in allowed_agent_ids: - agent = global_agent_registry.get_agent_by_id(agent_id) - if agent is not None: - models.append( - { - "model_name": f"a2a/{agent.agent_name}", - "litellm_params": { - "model": f"a2a/{agent.agent_name}", - "custom_llm_provider": "a2a", - }, - "model_info": { - "id": agent.agent_id, - "mode": "chat", - "db_model": True, - "created_by": agent.created_by, - "created_at": agent.created_at, - "updated_at": agent.updated_at, - }, - } - ) + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + for agent_id in allowed_agent_ids: + agent = global_agent_registry.get_agent_by_id(agent_id) + if agent is not None: + models.append( + { + "model_name": f"a2a/{agent.agent_name}", + "litellm_params": { + "model": f"a2a/{agent.agent_name}", + "custom_llm_provider": "a2a", + }, + "model_info": { + "id": agent.agent_id, + "mode": "chat", + "db_model": True, + "created_by": agent.created_by, + "created_at": agent.created_at, + "updated_at": agent.updated_at, + }, + } + ) + case _: + pass except Exception as e: verbose_proxy_logger.debug("Error appending agents to v2/model/info: %s", e) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3fba464dd23..d07ac0c5586 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,6 +13,7 @@ import asyncio import math import re import time +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast from fastapi import HTTPException, Request, status @@ -648,28 +649,29 @@ async def common_checks( ) async def _user_max_budget_check() -> None: - # 4.1 personal budget, if personal key - if ( - (team_object is None or team_object.team_id is None) - and user_object is not None - and user_object.max_budget is not None - ): - from litellm.proxy.proxy_server import get_current_spend + # 4.1 personal budget + if user_object is None or user_object.max_budget is None: + return + is_team_key: Final = team_object is not None and team_object.team_id is not None + if is_team_key and general_settings.get("apply_user_budget_to_team_keys") is not True: + return - user_budget: Final = user_object.max_budget - user_spend: Final = await get_current_spend( - counter_key=f"spend:user:{user_object.user_id}", - fallback_spend=user_object.spend or 0.0, + from litellm.proxy.proxy_server import get_current_spend + + user_budget: Final = user_object.max_budget + user_spend: Final = await get_current_spend( + counter_key=f"spend:user:{user_object.user_id}", + fallback_spend=user_object.spend or 0.0, + max_budget=user_budget, + ) + if math.isfinite(user_budget) and user_spend >= user_budget: + raise litellm.BudgetExceededError( + current_cost=user_spend, max_budget=user_budget, + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", + entity_type=Litellm_EntityType.USER.value, + entity_id=user_object.user_id, ) - if math.isfinite(user_budget) and user_spend >= user_budget: - raise litellm.BudgetExceededError( - current_cost=user_spend, - max_budget=user_budget, - message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}", - entity_type=Litellm_EntityType.USER.value, - entity_id=user_object.user_id, - ) # Each scope reads a distinct counter key with no cross-scope ordering # dependency, so the per-scope Redis-first reads run concurrently instead @@ -2833,7 +2835,7 @@ async def get_org_object( async def _get_resources_from_access_groups( - access_group_ids: list[str], + access_group_ids: Sequence[str], resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, @@ -2892,7 +2894,7 @@ async def _get_resources_from_access_groups( async def _get_models_from_access_groups( - access_group_ids: list[str], + access_group_ids: Sequence[str], prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, @@ -4262,6 +4264,10 @@ async def _project_soft_budget_check( ) +def _project_cache_key(project_id: str) -> str: + return f"project_id:{project_id}" + + async def get_project_object( project_id: str, prisma_client: PrismaClient | None, @@ -4279,7 +4285,7 @@ async def get_project_object( return None # Check cache first - cache_key: Final = f"project_id:{project_id}" + cache_key: Final = _project_cache_key(project_id) deserialized_project: Final = await user_api_key_cache.async_get_cache( key=cache_key, model_type=LiteLLM_ProjectTableCachedObj, @@ -4310,6 +4316,32 @@ async def get_project_object( return project_obj +async def delete_cached_project_object( + project_id: str, + user_api_key_cache: UserApiKeyCache, +) -> None: + """ + Every endpoint that mutates litellm_projecttable must call this: get_project_object + serves auth cache-first with no freshness check, so without invalidation a stale + project (e.g. a pre-update empty model allowlist) keeps being enforced until the + TTL expires (LIT-3803). Best-effort on both steps: the DB write has already + committed, so a cache backend error must not fail the endpoint; the stale entry + then expires via TTL. + """ + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation + + cache_key: Final = _project_cache_key(project_id) + try: + await user_api_key_cache.async_delete_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # best-effort eviction: any cache backend error must not fail the mutation + verbose_proxy_logger.warning( + "Failed to evict cached project entry %s; a stale project may be served until its TTL expires: %s", + cache_key, + e, + ) + await publish_auth_cache_invalidation(cache_key=cache_key) + + async def _organization_max_budget_check( valid_token: UserAPIKeyAuth | None, team_object: LiteLLM_TeamTable | None, diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index a6e5eb2a0a0..ff9211742f3 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -1,6 +1,7 @@ # What is this? ## Common checks for /v1/models and `/model/info` import copy +from collections.abc import Sequence from typing import Any, Final import litellm @@ -178,8 +179,8 @@ def get_team_models( def get_complete_model_list( - key_models: list[str], - team_models: list[str], + key_models: Sequence[str], + team_models: Sequence[str], proxy_model_list: list[str], user_model: str | None, infer_model_from_keys: bool | None, @@ -203,7 +204,7 @@ def get_complete_model_list( def append_unique(models): for model in models: - if model not in unique_models: + if model not in unique_models and model != SpecialModelNames.no_default_models.value: unique_models.append(model) if key_models: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9248576b599..4baa7b99a4f 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1044,6 +1044,22 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None: request.state.parent_otel_span = parent_otel_span +async def _read_request_body_deferring_parse_failure( + request: Request, +) -> tuple[dict, ProxyException | None]: + """Parse the body, returning a parse failure instead of raising it. + + A body that fails to parse is still a request from a known caller, so auth + must run (resolving identity onto the request's trace) before the 400 goes + out; the caller re-raises the returned exception once identity is seeded. + """ + try: + parsed_body: Final = await _read_request_body(request=request) + except ProxyException as parse_exception: + return {}, parse_exception # mutable-ok: request_data is a plain dict across the whole auth path + return populate_request_with_path_params(request_data=parsed_body, request=request), None + + async def _user_api_key_auth_builder( request: Request, api_key: str, @@ -2470,6 +2486,7 @@ async def _reserve_budget_after_common_checks( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, + apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True, fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True, ) @@ -2515,6 +2532,72 @@ def _resolve_request_principal(request: Request, valid_token: UserAPIKeyAuth) -> ) +async def _authorize_authenticated_request( + user_api_key_auth_obj: UserAPIKeyAuth, + request: Request, + request_data: dict, + route: str, + api_key: str, +) -> UserAPIKeyAuth | None: + """Authorize an already-authenticated request: disabled-route check, the single + ``common_checks`` gate (which also reserves budget), and end-user fallback + resolution. Returns the auth object the exception handler recovered when a check + failed but the request may proceed anyway, else ``None``. + """ + ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## + RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) + + # Single authorization point. Builder paths MUST NOT call common_checks. + # Route through the same exception handler the builder uses so + # authorization failures (ProxyException, or plain Exception from + # admin-only-route / model-access / budget checks) surface as + # ProxyException consistently with pre-refactor behavior. + try: + await _run_centralized_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request=request, + request_data=request_data, + route=route, + ) + except Exception as e: + return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( + e=e, + request=request, + request_data=request_data, + route=route, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + api_key=api_key, + resolved_identity=user_api_key_auth_obj, + ) + + # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return + # paths (no master key, /user/auth route, JWT short-circuits) that bypass + # the end-user resolution block. If those paths produced an auth obj + # without an ``end_user_id`` set, fall back to extracting from the request + # body so spend logs are still attributed correctly. Validation honours + # ``litellm.validate_end_user_id_in_db``. + if user_api_key_auth_obj.end_user_id is None: + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) + if raw_end_user_id is not None: + resolved_end_user_id: Final = await resolve_and_validate_end_user_id( + raw_end_user_id=raw_end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth_obj.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + route=route, + ) + if resolved_end_user_id is not None: + user_api_key_auth_obj.end_user_id = resolved_end_user_id + return None + + @tracer.wrap() async def user_api_key_auth( request: Request, @@ -2535,8 +2618,7 @@ async def user_api_key_auth( # close, and the trace never reaches the backend. _ensure_parent_otel_span_on_request_state(request) - request_data = await _read_request_body(request=request) - request_data = populate_request_with_path_params(request_data=request_data, request=request) + request_data, body_parse_exception = await _read_request_body_deferring_parse_failure(request=request) route: Final[str] = get_request_route(request=request) ## CHECK IF ROUTE IS ALLOWED @@ -2544,69 +2626,41 @@ async def user_api_key_auth( # triggers (key/user/team object reads) nest under it instead of flattening # onto the server span. No-op when OTel V2 isn't active. with phase_span(f"auth {route}"): - user_api_key_auth_obj: Final = await _user_api_key_auth_builder( - request=request, - api_key=api_key, - azure_api_key_header=azure_api_key_header, - anthropic_api_key_header=anthropic_api_key_header, - google_ai_studio_api_key_header=google_ai_studio_api_key_header, - azure_apim_header=azure_apim_header, - request_data=request_data, - custom_litellm_key_header=custom_litellm_key_header, - ) + try: + user_api_key_auth_obj: Final = await _user_api_key_auth_builder( + request=request, + api_key=api_key, + azure_api_key_header=azure_api_key_header, + anthropic_api_key_header=anthropic_api_key_header, + google_ai_studio_api_key_header=google_ai_studio_api_key_header, + azure_apim_header=azure_apim_header, + request_data=request_data, + custom_litellm_key_header=custom_litellm_key_header, + ) + except Exception: + # The body was read first, so a caller who sent both a malformed body and + # a rejected key used to get the 400; the response is unchanged, and the + # auth failure is still recorded on the trace by the handler that ran. + if body_parse_exception is not None: + raise body_parse_exception + raise user_api_key_auth_obj.budget_reservation = None - ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ## - RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj, request=request) - - # Single authorization point. Builder paths MUST NOT call common_checks. - # Route through the same exception handler the builder uses so - # authorization failures (ProxyException, or plain Exception from - # admin-only-route / model-access / budget checks) surface as - # ProxyException consistently with pre-refactor behavior. - try: - await _run_centralized_common_checks( + # A body that never parsed is authenticated (so the trace carries identity + # and this ``auth`` span) but not authorized: there is no model to check it + # against, and budget reservation would increment live spend counters that + # only the endpoint's post-call path releases; the endpoint never runs, since + # the parse failure is raised below. + if body_parse_exception is None: + recovered_auth_obj: Final = await _authorize_authenticated_request( user_api_key_auth_obj=user_api_key_auth_obj, request=request, request_data=request_data, route=route, - ) - except Exception as e: - return await UserAPIKeyAuthExceptionHandler._handle_authentication_error( - e=e, - request=request, - request_data=request_data, - route=route, - parent_otel_span=user_api_key_auth_obj.parent_otel_span, api_key=api_key, - resolved_identity=user_api_key_auth_obj, ) - - # Defense-in-depth: ``_user_api_key_auth_builder`` has multiple early-return - # paths (no master key, /user/auth route, JWT short-circuits) that bypass - # the end-user resolution block. If those paths produced an auth obj - # without an ``end_user_id`` set, fall back to extracting from the request - # body so spend logs are still attributed correctly. Validation honours - # ``litellm.validate_end_user_id_in_db``. - if user_api_key_auth_obj.end_user_id is None: - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) - - raw_end_user_id: Final = get_end_user_id_from_request_body(request_data, _safe_get_request_headers(request)) - if raw_end_user_id is not None: - resolved_end_user_id: Final = await resolve_and_validate_end_user_id( - raw_end_user_id=raw_end_user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth_obj.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - route=route, - ) - if resolved_end_user_id is not None: - user_api_key_auth_obj.end_user_id = resolved_end_user_id + if recovered_auth_obj is not None: + return recovered_auth_obj # Identity is now resolved. Seed it AFTER the auth span closes so the Baggage # persists on the request task (detaching the span's context token inside the @@ -2618,6 +2672,9 @@ async def user_api_key_auth( ) user_api_key_auth_obj.request_route = normalize_request_route(route) + if body_parse_exception is not None: + raise body_parse_exception + # Resolve caller identity once, here at the seam, into a single per-request # Principal projected off the key object the builder already fetched (no # second lookup). Downstream consumers read identity off this instead of diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index f7c332f2849..aef1c5ac17e 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -36,6 +36,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_original_file_id, prepare_data_with_credentials, update_batch_in_database, + validate_managed_id_requirement, ) from litellm.proxy.utils import handle_exception_on_proxy, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository @@ -176,6 +177,12 @@ async def create_batch( } input_file_id: Final = _create_batch_data.get("input_file_id", None) + await validate_managed_id_requirement( + resource_id=input_file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) unified_file_id: str | Literal[False] = False model_from_file_id = None @@ -392,6 +399,12 @@ async def retrieve_batch( data: dict = {} try: + await validate_managed_id_requirement( + resource_id=batch_id, + resource_kind="batch", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) model_from_id: Final = decode_model_from_file_id(batch_id) _retrieve_batch_request: Final = RetrieveBatchRequest( batch_id=batch_id, @@ -840,6 +853,13 @@ async def cancel_batch( data: dict = {} try: + await validate_managed_id_requirement( + resource_id=batch_id, + resource_kind="batch", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + # Check for encoded batch ID with model info model_from_id: Final = decode_model_from_file_id(batch_id) diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py new file mode 100644 index 00000000000..7fc8da42a3d --- /dev/null +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -0,0 +1,153 @@ +import asyncio +import json +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.config_sync_pubsub import ( + _ConfigSyncPubSub, + _pubsub_capable_client, + coordination_redis_cache, +) + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + +AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" +_POLL_TIMEOUT_SECONDS: Final = 1.0 +_BACKOFF_INITIAL_SECONDS: Final = 5.0 +_BACKOFF_MAX_SECONDS: Final = 60.0 + + +def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: + if redis_cache.namespace is None: + return AUTH_CACHE_INVALIDATION_CHANNEL + return f"{redis_cache.namespace}:{AUTH_CACHE_INVALIDATION_CHANNEL}" + + +@dataclass(frozen=True, slots=True) +class _CacheInvalidationMessage: + cache_key: str + + +def _cache_invalidation_message_json(cache_key: str) -> str: + return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key))) + + +def _cache_key_from_message_data(data: object) -> str | None: + if isinstance(data, bytes): + data = data.decode("utf-8", errors="replace") + if not isinstance(data, str): + return None + try: + parsed: Final = json.loads(data) + except json.JSONDecodeError: + return None + if not isinstance(parsed, dict): + return None + cache_key: Final = parsed.get("cache_key") + return cache_key if isinstance(cache_key, str) else None + + +async def publish_auth_cache_invalidation(cache_key: str) -> None: + """ + Best-effort broadcast so every worker drops its local in-memory copy of a + mutated management object; without this, only the handling worker and Redis + are evicted and other workers keep serving the stale object until its TTL. + """ + redis_cache: Final = coordination_redis_cache() + if redis_cache is None: + return + try: + client: Final = _pubsub_capable_client(redis_cache) + if client is None: + verbose_proxy_logger.debug( + "auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support", + cache_key, + ) + return + await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key)) + except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors + verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) + + +class AuthCacheInvalidationSubscriber: + __slots__ = ("_redis_cache", "_task", "_user_api_key_cache") + + def __init__( + self, + redis_cache: "RedisCache", + user_api_key_cache: "UserApiKeyCache", + ) -> None: + self._redis_cache = redis_cache + self._user_api_key_cache = user_api_key_cache + self._task: asyncio.Task[None] | None = None + + def start(self) -> None: + if self._task is not None: + return + self._task = asyncio.create_task(self._run()) + + async def stop(self) -> None: + task: Final = self._task + if task is None: + return + self._task = None + _ = task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def _run(self) -> None: + backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: exponential backoff accumulator across reconnects + while True: + try: + client = _pubsub_capable_client(self._redis_cache) # rebind-ok: re-resolved on every reconnect + if client is None: + verbose_proxy_logger.warning( + "auth cache invalidation subscriber disabled: cluster redis client has no pub/sub support; " + "cross-worker eviction falls back to the local cache TTL" + ) + return + pubsub = client.pubsub() # rebind-ok: fresh pubsub per reconnect + try: + await pubsub.subscribe(auth_cache_invalidation_channel(self._redis_cache)) + backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: reset after successful subscribe + await self._consume(pubsub) + finally: + await self._close_pubsub(pubsub) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect + verbose_proxy_logger.warning( + "auth cache invalidation subscriber redis error: %s; reconnecting in %.0fs", + e, + backoff_seconds, + ) + await asyncio.sleep(backoff_seconds) + backoff_seconds = min(backoff_seconds * 2, _BACKOFF_MAX_SECONDS) # rebind-ok: backoff accumulator + + async def _consume(self, pubsub: _ConfigSyncPubSub) -> None: + while True: + message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) + if message is None: + continue + self._apply_message(message) + + def _apply_message(self, message: object) -> None: + data: Final = message.get("data") if isinstance(message, dict) else None + cache_key: Final = _cache_key_from_message_data(data) + if cache_key is None: + return + in_memory_cache: Final = self._user_api_key_cache.in_memory_cache + if in_memory_cache is not None: + in_memory_cache.delete_cache(cache_key) + + @staticmethod + async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: + try: + await pubsub.aclose() + except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection + verbose_proxy_logger.debug("auth cache invalidation pubsub close failed: %s", e) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index a6141ac9217..3b3e9692eda 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -346,4 +346,4 @@ async def update_credential( return {"success": True, "message": "Credential updated successfully"} except Exception as e: - return handle_exception_on_proxy(e) + raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index da1652cdb61..b1f074c26b7 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -49,6 +49,7 @@ class AutoRouterTurnTransaction: cache_hit: bool cache_ttl_seconds: int | None cache_touched: bool + tier: str | None = None class TurnCacheFacts(NamedTuple): @@ -152,11 +153,13 @@ def build_autorouter_turn_transaction( return None usage_object_raw: Final = metadata.get("usage_object") cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None) + tier_raw: Final = routing_decision.get("tier") return AutoRouterTurnTransaction( api_key=api_key, session_id=_bounded_session_id(session_id), router_name=router_name, router_type=str(routing_decision.get("router_type") or "unknown"), + tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None, model=model, turn_at=turn_at, total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0), @@ -184,6 +187,8 @@ _COVERED: Final = _p("covered") _CACHE_HIT: Final = _p("cache_hit") _CACHE_TTL: Final = _p("cache_ttl_seconds") _TOUCHED: Final = _p("cache_touched") +_TIER: Final = f"{_p('tier')}::text" +_TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)" _IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at" _SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}" @@ -201,7 +206,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t ( last_model, models, turns, unordered_turns, covered_turns, cache_hits, same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, return_turns, return_hits, return_expired_misses, return_within_ttl_misses, - ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend + ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns ) VALUES ( {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, @@ -211,7 +216,8 @@ VALUES ( 0, 0, 0, 0, (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END), (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END), - {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8 + {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8, + {_TIER_DELTA} ) ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, @@ -242,6 +248,9 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET ELSE COALESCE((t.models -> {_MODEL} ->> 'ttl')::int, {_CACHE_TTL}::int) END) )), last_model = (CASE WHEN {_IN_ORDER} THEN {_MODEL} ELSE t.last_model END), + tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")} + THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1) + ELSE t.tier_turns END), first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at), last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at) """ diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 6863687081c..5f86490a474 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -486,6 +486,38 @@ class PrismaWrapper: os.environ[self._db_url_env_var] = _db_url return _db_url + @property + def engine_generation(self) -> int: + """How many query-engine replacements have completed on this wrapper. + + Bumped under `_reconnection_lock` only after a replacement engine has + connected, so a change across an await proves a *successful* planned + replacement happened in between — a replacement that failed (a real + outage) leaves it untouched. + """ + return self._engine_generation + + async def _reconnection_settled(self) -> None: + async with self._reconnection_lock: + pass + + async def wait_for_planned_engine_replacement(self, timeout_seconds: float) -> None: + """Wait, bounded, for an in-flight planned engine replacement to finish. + + Both replacement paths (`recreate_prisma_client` and + `_safe_refresh_token`) hold `_reconnection_lock` across their whole + kill/connect window, so re-acquiring it means the replacement has + settled one way or the other. Gives up silently on timeout: a caller + that stopped waiting must treat the replacement as not completed and + consult `engine_generation` rather than assume success. + """ + if timeout_seconds <= 0 or not self._reconnection_lock.locked(): + return + try: + await asyncio.wait_for(self._reconnection_settled(), timeout=timeout_seconds) + except asyncio.TimeoutError: + return + async def recreate_prisma_client( self, new_db_url: str, diff --git a/litellm/proxy/fine_tuning_endpoints/endpoints.py b/litellm/proxy/fine_tuning_endpoints/endpoints.py index f8ffb77edb8..a13ad00713d 100644 --- a/litellm/proxy/fine_tuning_endpoints/endpoints.py +++ b/litellm/proxy/fine_tuning_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, + validate_managed_id_requirement, ) from litellm.proxy.utils import handle_exception_on_proxy from litellm.types.utils import LiteLLMFineTuningJob @@ -134,6 +135,18 @@ async def create_fine_tuning_job( ## CHECK IF MANAGED FILE ID unified_file_id: str | Literal[False] = False training_file: Final = fine_tuning_request.training_file + await validate_managed_id_requirement( + resource_id=training_file, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + await validate_managed_id_requirement( + resource_id=fine_tuning_request.validation_file, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) response: LiteLLMFineTuningJob | None = None if training_file: unified_file_id = _is_base64_encoded_unified_file_id(training_file) @@ -246,6 +259,12 @@ async def retrieve_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") + await validate_managed_id_requirement( + resource_id=fine_tuning_job_id, + resource_kind="fine-tuning job", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( @@ -513,6 +532,12 @@ async def cancel_fine_tuning_job( try: if premium_user is not True: raise ValueError(f"Only premium users can use this endpoint + {CommonProxyErrors.not_premium_user.value}") + await validate_managed_id_requirement( + resource_id=fine_tuning_job_id, + resource_kind="fine-tuning job", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 761d8aabc8a..b68d4a68b79 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -6,10 +6,10 @@ import concurrent.futures import inspect import json import os -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timezone from types import UnionType -from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, Union, cast, get_args, get_origin +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, Union, cast, get_args, get_origin from urllib.parse import urlparse from fastapi import APIRouter, Depends, HTTPException, Request @@ -54,8 +54,8 @@ from litellm.types.guardrails import ( if TYPE_CHECKING: from types import CodeType - from prisma.actions import LiteLLM_GuardrailsTableActions from prisma.models import LiteLLM_GuardrailsTable + from pydantic.fields import FieldInfo from litellm.proxy.utils import PrismaClient @@ -65,24 +65,44 @@ router: Final = APIRouter() GUARDRAIL_REGISTRY: Final = GuardrailRegistry() -def _guardrails_table(prisma_client: "PrismaClient") -> "LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]": - table: Final[LiteLLM_GuardrailsTableActions[LiteLLM_GuardrailsTable]] = GuardrailsRepository(prisma_client).table +class _GuardrailsTableActions(Protocol): + async def create(self, data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": ... + + async def delete(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ... + + async def find_unique(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ... + + async def find_many( + self, where: Mapping[str, object], order: Mapping[str, str] + ) -> "Sequence[LiteLLM_GuardrailsTable]": ... + + async def update( + self, where: Mapping[str, object], data: Mapping[str, object] + ) -> "LiteLLM_GuardrailsTable | None": ... + + +def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]: + return mapping + + +def _guardrails_table(prisma_client: "PrismaClient") -> _GuardrailsTableActions: + table: Final[_GuardrailsTableActions] = GuardrailsRepository(prisma_client).table return table async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": - row: Final[LiteLLM_GuardrailsTable] = await GuardrailsRepository(prisma_client).table.create(data=data) + row: Final = await _guardrails_table(prisma_client).create(data=data) return row async def _delete_guardrail_row(prisma_client: "PrismaClient", where: Mapping[str, object]) -> None: - await GuardrailsRepository(prisma_client).table.delete(where=where) + await _guardrails_table(prisma_client).delete(where=where) async def _find_team_guardrail_rows( prisma_client: "PrismaClient", where: Mapping[str, object] ) -> "Sequence[LiteLLM_GuardrailsTable]": - rows: Final[Sequence[LiteLLM_GuardrailsTable]] = await GuardrailsRepository(prisma_client).table.find_many( + rows: Final = await _guardrails_table(prisma_client).find_many( where=where, order={"created_at": "desc"}, ) @@ -499,10 +519,12 @@ async def update_guardrail( if existing_guardrail is None: raise HTTPException(status_code=404, detail=f"Guardrail with ID {guardrail_id} not found") - result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db( - guardrail_id=guardrail_id, - guardrail=request.guardrail, - prisma_client=prisma_client, + result: Final = _as_str_object_mapping( + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=request.guardrail, + prisma_client=prisma_client, + ) ) guardrail_name: Final = result.get("guardrail_name", "Unknown") @@ -613,7 +635,7 @@ class RegisterGuardrailRequest(BaseModel): """Request body for POST /guardrails/register. Follows Generic Guardrail API config.""" guardrail_name: str - litellm_params: dict[str, Any] # guardrail, mode, api_base required; api_key, headers, etc. optional + litellm_params: dict[str, object] # guardrail, mode, api_base required; api_key, headers, etc. optional guardrail_info: dict[str, object] | None = None team_id: str | None = None @@ -1172,12 +1194,14 @@ async def patch_guardrail( ) # Update litellm_params if default_on is provided or pii_entities_config is provided - litellm_params = LitellmParams(**dict(existing_guardrail.get("litellm_params", {}))) + existing_litellm_params: Final = _as_str_object_mapping(dict(existing_guardrail.get("litellm_params", {}))) + litellm_params = LitellmParams(**existing_litellm_params) if request.litellm_params is not None: requested_litellm_params: Final = request.litellm_params.model_dump(exclude_unset=True) litellm_params_dict: Final = litellm_params.model_dump(exclude_unset=True) litellm_params_dict.update(requested_litellm_params) - litellm_params = LitellmParams(**litellm_params_dict) + merged_litellm_params: Final = _as_str_object_mapping(litellm_params_dict) + litellm_params = LitellmParams(**merged_litellm_params) # Update guardrail_info if provided guardrail_info: Final = ( @@ -1193,10 +1217,12 @@ async def patch_guardrail( litellm_params=litellm_params, guardrail_info=guardrail_info, ) - result: Final = await GUARDRAIL_REGISTRY.update_guardrail_in_db( - guardrail_id=guardrail_id, - guardrail=guardrail, - prisma_client=prisma_client, + result: Final = _as_str_object_mapping( + await GUARDRAIL_REGISTRY.update_guardrail_in_db( + guardrail_id=guardrail_id, + guardrail=guardrail, + prisma_client=prisma_client, + ) ) guardrail_name = result.get("guardrail_name", "Unknown") @@ -1552,31 +1578,46 @@ async def validate_blocked_words_file(request: dict[str, str]): return {"valid": False, "error": f"Validation error: {e}"} -def _get_field_type_from_annotation(field_annotation: Any) -> str: +def _dunder_origin(annotation: object) -> object: + origin: Final[object] = getattr(annotation, "__origin__", None) + return origin + + +def _dunder_name(annotation: object) -> object: + name: Final[object] = getattr(annotation, "__name__", None) + return name + + +def _dunder_args(annotation: object) -> tuple[object, ...]: + args: Final[tuple[object, ...]] = getattr(annotation, "__args__", ()) + return args + + +def _get_field_type_from_annotation(field_annotation: object) -> str: """ Convert a Python type annotation to a UI-friendly type string """ # Handle Union types (like Optional[T]) if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType: # For Optional[T], get the non-None type - args: Final = get_args(field_annotation) + args: Final[tuple[object, ...]] = get_args(field_annotation) non_none_args: Final = [arg for arg in args if arg is not type(None)] if non_none_args: field_annotation = non_none_args[0] # Handle List types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is list: + if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is list: return "array" # Handle Dict types - if hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is dict: + if hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is dict: return "dict" # Handle Literal types if hasattr(field_annotation, "__origin__") and hasattr(field_annotation, "__args__"): # Check for Literal types (Python 3.8+) - origin: Final = field_annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": + origin: Final = _dunder_origin(field_annotation) + if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal": return "select" # For dropdown/select inputs # Handle basic types @@ -1595,66 +1636,66 @@ def _get_field_type_from_annotation(field_annotation: Any) -> str: return "string" -def _extract_literal_values(annotation: Any) -> list[str]: +def _extract_literal_values(annotation: object) -> Sequence[object]: """ Extract literal values from a Literal type annotation """ if hasattr(annotation, "__origin__") and hasattr(annotation, "__args__"): - origin: Final = annotation.__origin__ - if hasattr(origin, "__name__") and origin.__name__ == "Literal": - return list(annotation.__args__) + origin: Final = _dunder_origin(annotation) + if hasattr(origin, "__name__") and _dunder_name(origin) == "Literal": + return list(_dunder_args(annotation)) return [] -def _get_dict_key_options(field_annotation: Any) -> list[str] | None: +def _get_dict_key_options(field_annotation: object) -> Sequence[object] | None: """ Extract key options from Dict[Literal[...], T] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict + and _dunder_origin(field_annotation) is dict and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 2: key_type: Final = args[0] return _extract_literal_values(key_type) return None -def _get_dict_value_type(field_annotation: Any) -> str: +def _get_dict_value_type(field_annotation: object) -> str: """ Get the value type from Dict[K, V] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is dict + and _dunder_origin(field_annotation) is dict and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 2: value_type: Final = args[1] return _get_field_type_from_annotation(value_type) return "string" -def _get_list_element_options(field_annotation: Any) -> list[str] | None: +def _get_list_element_options(field_annotation: object) -> Sequence[object] | None: """ Extract element options from List[Literal[...]] types """ if ( hasattr(field_annotation, "__origin__") - and field_annotation.__origin__ is list + and _dunder_origin(field_annotation) is list and hasattr(field_annotation, "__args__") ): - args: Final = field_annotation.__args__ + args: Final = _dunder_args(field_annotation) if len(args) >= 1: element_type: Final = args[0] return _extract_literal_values(element_type) return None -def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool: +def _should_skip_optional_params(field_name: str, field_annotation: object) -> bool: """Check if optional_params field should be skipped (not meaningfully overridden).""" if field_name != "optional_params": return False @@ -1664,12 +1705,12 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Check if the annotation is still a generic TypeVar (not specialized) if isinstance(field_annotation, TypeVar) or ( - hasattr(field_annotation, "__origin__") and field_annotation.__origin__ is TypeVar + hasattr(field_annotation, "__origin__") and _dunder_origin(field_annotation) is TypeVar ): return True # Also skip if it's a generic type that wasn't specialized - if hasattr(field_annotation, "__name__") and field_annotation.__name__ in ( + if hasattr(field_annotation, "__name__") and _dunder_name(field_annotation) in ( "T", "TypeVar", ): @@ -1677,18 +1718,18 @@ def _should_skip_optional_params(field_name: str, field_annotation: Any) -> bool # Handle Optional[T] where T is still a TypeVar if hasattr(field_annotation, "__args__"): - non_none_args: Final = [arg for arg in field_annotation.__args__ if arg is not type(None)] + non_none_args: Final = [arg for arg in _dunder_args(field_annotation) if arg is not type(None)] if non_none_args and isinstance(non_none_args[0], TypeVar): return True return False -def _unwrap_optional_type(field_annotation: Any) -> Any: +def _unwrap_optional_type(field_annotation: object) -> object: """Unwrap Optional types to get the actual type.""" if get_origin(field_annotation) is Union or get_origin(field_annotation) is UnionType: # For Optional[BaseModel], get the non-None type - args: Final = get_args(field_annotation) + args: Final[tuple[object, ...]] = get_args(field_annotation) non_none_args: Final = [arg for arg in args if arg is not type(None)] if non_none_args: return non_none_args[0] @@ -1696,20 +1737,20 @@ def _unwrap_optional_type(field_annotation: Any) -> Any: def _build_field_dict( - field: Any, - field_annotation: Any, + field: "FieldInfo", + field_annotation: object, description: str, required: bool, -) -> dict[str, Any]: +) -> dict[str, object]: """Build field dictionary for non-nested fields.""" # Determine the field type from annotation field_type = _get_field_type_from_annotation(field_annotation) # Check for custom UI type override - field_json_schema_extra: Final = getattr(field, "json_schema_extra", {}) + field_json_schema_extra: Final[Mapping[str, object]] = getattr(field, "json_schema_extra", {}) if field_json_schema_extra and "ui_type" in field_json_schema_extra: ui_type: Final = field_json_schema_extra["ui_type"] - field_type = ui_type.value if hasattr(ui_type, "value") else ui_type + field_type = getattr(ui_type, "value", ui_type) elif field_json_schema_extra and "type" in field_json_schema_extra: field_type = field_json_schema_extra["type"] @@ -1748,8 +1789,9 @@ def _build_field_dict( field_dict["options"] = literal_options # Add default value if it exists - if field.default is not None and field.default is not ...: - field_dict["default_value"] = field.default + field_default: Final[object] = getattr(field, "default", None) + if field_default is not None and field_default is not ...: + field_dict["default_value"] = field_default # Copy min, max, step from json_schema_extra for number/percentage inputs if field_json_schema_extra: @@ -1763,7 +1805,7 @@ def _build_field_dict( def _extract_fields_recursive( model: type[BaseModel], depth: int = 0, -) -> dict[str, Any]: +) -> dict[str, object]: # Check if we've exceeded the maximum recursion depth if depth > DEFAULT_MAX_RECURSE_DEPTH: raise HTTPException( @@ -1817,7 +1859,7 @@ def _extract_fields_recursive( return fields -def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, Any]: +def _get_fields_from_model(model_class: type[BaseModel]) -> dict[str, object]: """ Get the fields from a Pydantic model as a nested dictionary structure """ @@ -2141,7 +2183,26 @@ def _resolve_guardrail_input_type(active_guardrail: CustomGuardrail, input_type: return "response" if input_type == "response" else "request" -def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGuardrailRequest) -> None: +class _GuardrailLoggingObj(Protocol): + call_type: str + model_call_details: dict[str, object] + + @property + def update_messages(self) -> "Callable[..., object]": ... + + @property + def async_success_handler(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def success_handler(self) -> "Callable[..., object]": ... + + +class _GuardrailProxyLogging(Protocol): + @property + def post_call_success_hook(self) -> "Callable[..., Awaitable[object]]": ... + + +def _patch_logging_obj_for_guardrail(litellm_logging_obj: _GuardrailLoggingObj, request: ApplyGuardrailRequest) -> None: """Configure the logging object so Langfuse/OTEL extract input and output correctly.""" litellm_logging_obj.call_type = "pass_through_endpoint" litellm_logging_obj.model_call_details["call_type"] = "pass_through_endpoint" @@ -2151,8 +2212,8 @@ def _patch_logging_obj_for_guardrail(litellm_logging_obj: Any, request: ApplyGua async def _emit_guardrail_success_logs( - proxy_logging_obj: Any, - litellm_logging_obj: Any, + proxy_logging_obj: _GuardrailProxyLogging, + litellm_logging_obj: _GuardrailLoggingObj | None, data: dict, user_api_key_dict: UserAPIKeyAuth, response: ApplyGuardrailResponse, diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index e9e729fb118..eecbce57468 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -9,11 +9,14 @@ import os import sys sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path +import asyncio import copy import json +import re import sys -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Mapping, Sequence from datetime import datetime, timezone +from itertools import accumulate, groupby from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast import httpx @@ -23,6 +26,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys @@ -46,6 +50,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailOutput, BedrockGuardrailQualifier, BedrockGuardrailResponse, + BedrockGuardrailUsage, BedrockRequest, BedrockTextContent, ) @@ -53,6 +58,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from botocore.awsrequest import AWSPreparedRequest + from botocore.credentials import Credentials from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -71,6 +77,17 @@ from litellm.types.utils import ( GUARDRAIL_NAME: Final = "bedrock" _BEDROCK_DYNAMIC_BODY_DENYLIST: Final = frozenset({"content", "source"}) +_BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS: Final = ( + "text unit", + "maximum input size", + "content size", + "too long", + "too large", + "exceeds the maximum", +) +_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES: Final = 3 +_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS: Final = 0.5 +_BEDROCK_WHITESPACE: Final = re.compile(r"\s") # Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required). _BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH: Final = "/guardrail-checks/invoke" # InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with @@ -118,6 +135,29 @@ class GuardrailMessageFilterResult(NamedTuple): target_indices: list[int] | None +class BedrockContentChunkResult(NamedTuple): + """One chunk's ApplyGuardrail response, paired with enough bookkeeping to + reconstruct global masked-output positions once every chunk is back. + + `content` is the exact content items this chunk was called with -- needed + so an all-clear chunk (empty `outputs`) can still contribute one unmasked + placeholder per item it covers, keeping every later chunk's masked text + aligned to its original global position. `fragment_group_size` is 1 for an + ordinary chunk, and otherwise the total number of consecutive chunk results + that together make up ONE original content item's own text (split because a + list of length 1 could not be bisected by list length). All of them must be + concatenated back into that one item's masked output rather than treated as + separate items. It is a count rather than a boolean because one item can be + bisected more than once: two levels of splitting produce four fragments for + a single item, not two, and grouping them in fixed pairs would emit two + outputs for one message and shift every later message's masked text. + """ + + response: BedrockGuardrailResponse + content: tuple[BedrockContentItem, ...] + fragment_group_size: int + + class ApplyGuardrailMessageSelection(NamedTuple): """Messages selected for an apply_guardrail scan + write-back metadata.""" @@ -168,12 +208,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): content_filter_threshold: float | None = 0.5, prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, + chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.guardrailIdentifier = guardrailIdentifier self.guardrailVersion = guardrailVersion self.guardrail_provider = "bedrock" + self.chunk_budget_chars = chunk_budget_chars self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only")) # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks` @@ -759,12 +801,35 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None = None, logging_event_type: GuardrailEventHooks | None = None, ) -> BedrockGuardrailResponse: + """Scan `messages`/`response` with ApplyGuardrail, chunking if it is too large. + + Content is bin-packed into budget-sized batches and each batch posted + sequentially, every batch independently falling back to bisection if AWS + rejects it. The per-batch responses are merged so callers cannot tell whether + chunking happened. + + Content using contextual grounding opts out of chunking entirely: grounding is + scored holistically against the whole reference source, so bisecting it would + fragment that evaluation and yield misleading scores. Such a request keeps the + old behavior of surfacing a too-large error rather than being split. + + `logging_event_type` drives what UI and spend logs report. It is distinct from + Bedrock's `source`, which is INPUT vs OUTPUT for the API body and must not be + confused with the proxy hook (pre_call / during_call / post_call); when omitted, + the legacy source-derived mapping is kept for backward compatibility. + + A guardrail *block* is logged where it happens, in + `_post_apply_guardrail_content`, because chunking stops immediately and there is + no later merged response to log instead. Everything else that fails out of the + chunking flow (an unrecoverable too-large error, a non-size validation error, + exhausted throttle retries) is a genuine end-to-end failure of this one logical + guardrail call and is logged exactly once here. + """ start_time: Final = datetime.now(timezone.utc) credentials, aws_region_name = self._load_credentials() bedrock_request_data: Final[dict] = dict( self.convert_to_bedrock_format(source=source, messages=messages, response=response) ) - bedrock_guardrail_response: BedrockGuardrailResponse = BedrockGuardrailResponse() api_key: str | None = None if request_data: dynamic_request_body_params = self.get_guardrail_dynamic_request_body_params(request_data=request_data) @@ -778,6 +843,257 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if request_data.get("api_key") is not None: api_key = request_data["api_key"] + event_type: Final = ( + logging_event_type + if logging_event_type is not None + else (GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call) + ) + + content: Final[tuple[BedrockContentItem, ...]] = tuple(bedrock_request_data.get("content") or ()) + allow_chunking: Final = not self._content_uses_contextual_grounding(content) + + try: + responses: Final = await self._apply_guardrail_content_with_chunking( + content=content, + base_request_data=bedrock_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + except HTTPException as exc: + if not isinstance(exc.detail, dict): + self._log_apply_guardrail_failure( + detail=exc.detail, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + raise + merged_response: Final = self._merge_bedrock_guardrail_responses(responses) + self._log_apply_guardrail_success( + merged_response=merged_response, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + return merged_response + + async def _apply_guardrail_content_with_chunking( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + allow_chunking: bool, + ) -> tuple[BedrockContentChunkResult, ...]: + """Post `content` to ApplyGuardrail, chunking only if AWS rejects it as too large. + + Tries `content` as a single call first. AWS's per-request "maximum input + size in text units" quota is account/region/policy-dependent and cannot be + predicted ahead of time, so it is only ever discovered reactively: on an + error whose message indicates the input was too large (a ThrottlingException + in practice, a ValidationException per the docs -- see + ``_is_input_too_large_error``), the content is re-sent in smaller pieces. + + Probing with the whole payload first is what keeps a request AWS would have + accepted at exactly one call. Packing into fixed batches up front instead + would split conversations AWS was happy to take whole, multiplying billed + calls and guardrail latency on traffic that never had a size problem, and + no fixed budget can avoid that because the real cap is unknown here. + + Once a rejection proves the payload is over the cap, a multi-item payload is + re-sent as ``chunk_budget_chars``-sized batches rather than bisected: that + reaches a working size in one step instead of paying an O(log n) ladder of + rejected calls. Bisection remains the fallback for anything bin-packing + cannot make smaller, which is what makes the recursion terminate: a batch + already inside the budget packs back to itself, so it falls through to the + split below. A single oversized + content item (one very long message) is split by its own text instead of + by list length, since a list of length 1 has no items left to bisect -- + the resulting fragments all carry a ``fragment_group_size`` so the merge + step can recombine them into the one content item they came from, rather + than treating each fragment as its own item when reconstructing positions + for masking. That count covers however many fragments the item ended up + split into, not just two, since it can be bisected repeatedly: the + outermost single-item split stamps the total leaf count on every leaf + below it, overwriting any smaller count an inner split had set. A real + guardrail block on any (sub-)chunk raises immediately + -- callers must not lose that signal by continuing to post the remaining + chunks. + """ + try: + response: Final = await self._post_apply_guardrail_content_with_retry( + content=content, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + return ( + BedrockContentChunkResult( + response=response, + content=tuple(content), + fragment_group_size=1, + ), + ) + except HTTPException as exc: + if allow_chunking and self._is_input_too_large_error(exc.detail): + batches: Final = self._bin_pack_bedrock_content(content, budget=self.chunk_budget_chars) + if len(batches) > 1: + verbose_proxy_logger.warning( + "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; " + "re-sending as %d batches of at most %d characters", + len(content), + len(batches), + self.chunk_budget_chars, + ) + batch_results: Final = [ # mutable-ok: await needs a list comprehension; frozen to a tuple below + await self._apply_guardrail_content_with_chunking( + content=batch, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + for batch in batches + ] + return tuple(result for results in batch_results for result in results) + split_content: Final = self._split_bedrock_content(content) + if split_content is None: + raise + first_half, second_half = split_content + is_single_item_text_split: Final = len(content) == 1 + verbose_proxy_logger.warning( + "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; " + "splitting into %d + %d and retrying each", + len(content), + len(first_half), + len(second_half), + ) + first_results: Final = await self._apply_guardrail_content_with_chunking( + content=first_half, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + second_results: Final = await self._apply_guardrail_content_with_chunking( + content=second_half, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + combined_results: Final = tuple(first_results) + tuple(second_results) + if is_single_item_text_split: + return tuple( + result._replace(fragment_group_size=len(combined_results)) for result in combined_results + ) + return combined_results + raise + + async def _post_apply_guardrail_content_with_retry( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> BedrockGuardrailResponse: + """Post one ApplyGuardrail call for `content`, retrying with exponential + backoff on AWS ThrottlingException (HTTP 429). + + Chunking already trades one oversized call for several smaller ones, so + retries here are capped low -- they must not multiply per-request latency + by an order of magnitude when the account's per-second text-unit quota is + the binding constraint rather than the per-request size quota. + + A too-large rejection is deliberately excluded from the retry. AWS reports + it as a ThrottlingException (429), not only as a ValidationException, but + unlike a genuine throttle it is not transient: re-posting the same + oversized content can never succeed. Retrying it would burn every backoff + sleep and every (billed) attempt before the caller's bisection gets a + chance to split the content, at every level of the recursion. + """ + for attempt in range(_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + 1): + try: + return await self._post_apply_guardrail_content( + content=content, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + except HTTPException as exc: + if ( + exc.status_code != 429 + or self._is_input_too_large_error(exc.detail) + or attempt >= _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + ): + raise + await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2**attempt)) + raise HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted") + + async def _post_apply_guardrail_content( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> BedrockGuardrailResponse: + """Make exactly one signed ApplyGuardrail HTTP call for `content` and + parse the result. Raises HTTPException on a guardrail block or any + non-200 response (including 429, handled by the retry wrapper above). + + AWS also reports some failures inside a 200 body, tagging ``Output.__type`` + with an Exception marker. Those deliberately do NOT raise: the request proceeds, + matching the behaviour of this code before chunking existed. The marker survives + the merge, so the one consolidated log entry still records + ``guardrail_failed_to_respond`` rather than a success. Making that path fail + closed is a separate change, tracked apart from this PR, and belongs behind the + existing ``unreachable_fallback`` setting rather than a hardcoded status. + + A block is logged here rather than by the caller: it ends the whole chunking + flow immediately, with no further chunks attempted, so there is no later + merged response for the caller to log instead. + """ + bedrock_request_data: Final = { # mutable-ok: outbound JSON request body + **base_request_data, + "content": content, + } # mutable-ok: outbound JSON request body prepared_request: Final = self._prepare_request( credentials=credentials, data=bedrock_request_data, @@ -792,42 +1108,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) - # UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API - # body, which must not be confused with the proxy hook (pre_call / during_call / - # post_call). When omitted, keep legacy mapping for backward compatibility. - if logging_event_type is not None: - event_type = logging_event_type - else: - event_type = GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call - httpx_response: Final = await self._sign_and_post( prepared_request=prepared_request, request_data=request_data, event_type=event_type, start_time=start_time, + log_transport_failure=False, ) - ######################################################### - # Add guardrail information to request trace - ######################################################### - _json_response: Final = httpx_response.json() - tracing_detail: Final = self._build_tracing_detail(_json_response) - - # Raw Bedrock JSON is passed here; match/regex redaction runs once inside - # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. - self.add_standard_logging_guardrail_information_to_request_data( - guardrail_provider=self.guardrail_provider, - guardrail_json_response=_json_response, - request_data=request_data or {}, - guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response), - start_time=start_time.timestamp(), - end_time=datetime.now(timezone.utc).timestamp(), - duration=(datetime.now(timezone.utc) - start_time).total_seconds(), - event_type=event_type, - tracing_detail=tracing_detail or None, - ) - ######################################################### if httpx_response.status_code == 200: + _json_response: Final = httpx_response.json() # check if the response was flagged verbose_proxy_logger.debug( "Bedrock AI response : %s", @@ -835,19 +1125,462 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): + self._log_apply_guardrail_attempt( + httpx_response=httpx_response, + json_response=_json_response, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) raise self._get_http_exception_for_blocked_guardrail( bedrock_guardrail_response, request_data=request_data ) - else: - status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response) - verbose_proxy_logger.error( - "Bedrock AI: error in response. Status code: %s, response: %s", - httpx_response.status_code, - httpx_response.text, - ) - raise HTTPException(status_code=status_code, detail=detail_message) + return bedrock_guardrail_response - return bedrock_guardrail_response + status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response) + verbose_proxy_logger.error( + "Bedrock AI: error in response. Status code: %s, response: %s", + httpx_response.status_code, + httpx_response.text, + ) + raise HTTPException(status_code=status_code, detail=detail_message) + + def _log_apply_guardrail_attempt( + self, + httpx_response: httpx.Response, + json_response: dict, # mutable-ok: raw AWS JSON payload + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log a single ApplyGuardrail HTTP attempt as-is (its own status, + derived from its own response). Used only for the blocked-content + case, which ends the whole chunking flow immediately.""" + tracing_detail: Final = self._build_tracing_detail(BedrockGuardrailResponse(**json_response)) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response=json_response, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response), + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + tracing_detail=tracing_detail or None, + ) + + def _log_apply_guardrail_success( + self, + merged_response: BedrockGuardrailResponse, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log one logical ApplyGuardrail call -- possibly several chunk calls + under the hood -- using its final merged response, so a chunked + request produces exactly one telemetry entry, the same as an + unchunked one would. + + AWS can report a failure inside an HTTP 200 body by tagging + ``Output.__type`` with an exception marker. That marker survives the merge, + so the status is derived from the merged response rather than assumed to be + a success, which is what the pre-chunking code reported for that shape.""" + tracing_detail: Final = self._build_tracing_detail(merged_response) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response=dict(merged_response), # mutable-ok: logging helper requires a dict + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status=( + "guardrail_failed_to_respond" + if "Exception" in str((merged_response.get("Output") or {}).get("__type", "")) + else "success" + ), + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + tracing_detail=tracing_detail or None, + ) + + def _log_apply_guardrail_failure( + self, + detail: object, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log one logical ApplyGuardrail call that failed end-to-end (an + unrecoverable too-large error, a non-size validation error, or + exhausted throttle retries) as a single failure, rather than logging + every failed attempt chunking made along the way.""" + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={"error": str(detail)}, # mutable-ok: logging helper requires a dict + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + ) + + @staticmethod + def _content_uses_contextual_grounding(content: Sequence[BedrockContentItem]) -> bool: + """True if any content item carries a contextual-grounding qualifier + (``grounding_source``, ``query``, or the ``guard_content`` the response + itself is tagged with once grounding is present).""" + for item in content: + if (item.get("text") or {}).get("qualifiers"): # mutable-ok: read-only empty fallback + return True + return False + + @staticmethod + def _bin_pack_bedrock_content( + content: Sequence[BedrockContentItem], + budget: int, + ) -> tuple[tuple[BedrockContentItem, ...], ...]: + """Pack whole content items, in order, into batches whose combined text + length stays within `budget`, in a single pass that carries the running + total rather than re-summing the open batch per item. + + This is the fast-path half of the hybrid chunking strategy: bin-packing + at a conservative fixed budget keeps the common case at O(n / budget) + ApplyGuardrail calls instead of the O(log n) round trips pure reactive + bisection pays on every oversized request. An item whose own text + already exceeds `budget` is not split here -- it becomes its own + (still oversized) batch and is sent as-is; if AWS rejects that batch as + too large, `_apply_guardrail_content_with_chunking`'s existing + recursive-bisection fallback takes over for that batch only. + + `budget` comes from the guardrail's ``chunk_budget_chars`` setting and + defaults to 25,000, matching ApplyGuardrail's default quota of 25 text + units (roughly 1,000 characters each) per second. Packing to that size and + posting sequentially is what keeps chunking from tripping the rate quota + and trading a size error for a throttle. Accounts with raised quotas can + configure a larger budget to spend fewer calls. + + The budget is not a correctness dependency either way. AWS's effective cap + varies by account, region, and policy, is not a fixed character count, and + cannot be read from config, so any batch it still rejects falls back to + bisection, which self-corrects however wrong the value was. An over-large + budget therefore costs one extra probe-and-bisect round trip rather than + failing the request. + """ + if not content: + return (tuple(content),) + + lengths: Final = tuple(len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content) + + def assign(carried: tuple[int, int], length: int) -> tuple[int, int]: + batch_index, used = carried + if used + length <= budget: + return batch_index, used + length + return batch_index + 1, length + + batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:]) + return tuple( + tuple(item for _, item in group) + for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0]) + ) + + @staticmethod + def _split_bedrock_content( + content: Sequence[BedrockContentItem], + ) -> tuple[tuple[BedrockContentItem, ...], tuple[BedrockContentItem, ...]] | None: + """Bisect `content` into two roughly-equal, non-empty halves. + + When `content` already holds more than one item, it is split by list + length. When it holds exactly one item, that item's own text is split + instead (a list of length 1 has no items left to bisect, but one very + long message is still a single content item) -- at the whitespace + character nearest the midpoint rather than a raw character index, so + the cut never lands inside a word/token. This is a plain, lossless + cut with no overlap: concatenating the two fragments in order always + reproduces the original text exactly, so merging back at + ``_merge_logical_unit_outputs`` needs no reconciliation step. + + Known, accepted limitation: whitespace splitting only guards against + *accidentally* severing a single token (one denied word, one PII + pattern) across the cut. It does not, and cannot without an overlap + window, stop a *multi-word* denied phrase deliberately positioned to + straddle the boundary -- each fragment can scan clean on its own and + still reassemble into the flagged phrase. AWS's own guidance on this + API acknowledges the same gap for input chunking ("a critical piece of + text could span two (or more) chunks if not carefully divided") with + no documented resolution, and overlap-and-reconcile was evaluated and + rejected for this PR: AWS's masking output has no documented + length-preservation guarantee, so reconciling an overlap region against + masked text is not sound in general. Out of scope for this PR. + + Returns None when there is nothing left to split -- a single item + whose text is too short to halve into two non-empty pieces -- so the + caller can give up and propagate the original too-large error instead + of recursing forever. + """ + if len(content) > 1: + midpoint: Final = max(1, len(content) // 2) + return tuple(content[:midpoint]), tuple(content[midpoint:]) + + text_content: Final = content[0].get("text") or BedrockTextContent() + text: Final = text_content.get("text") or "" + if len(text) < 2: + return None + split_at: Final = BedrockGuardrail._nearest_whitespace_split_index(text) + qualifiers: Final = text_content.get("qualifiers") + + def fragment(piece: str) -> BedrockContentItem: + block: Final = ( + BedrockTextContent(text=piece, qualifiers=qualifiers) if qualifiers else BedrockTextContent(text=piece) + ) + return BedrockContentItem(text=block) + + return (fragment(text[:split_at]),), (fragment(text[split_at:]),) + + @staticmethod + def _nearest_whitespace_split_index(text: str) -> int: + """Return the index nearest `text`'s midpoint that falls on a whitespace + boundary, so splitting `text[:i]` / `text[i:]` there never severs a word. + + Any Unicode whitespace counts, not just an ASCII space. Matching only `" "` + would leave the boundary unguarded for exactly the payloads that get large + enough to need splitting: JSON lines, source code, logs and transcripts are + newline or tab delimited, so a deny-listed word sitting at the midpoint of + one would be cut in half, scan clean on both fragments, and reassemble + intact. + + The returned index always leaves both sides non-empty, which is what makes + the caller's recursion terminate. A boundary that would put the split at 0 + or at ``len(text)`` is discarded: it would hand back a fragment identical to + the text just rejected as too large, AWS would reject that again, and each + retry would re-split it into the same unchanged fragment until the stack ran + out. The dangerous shape is a text whose only space at or after the midpoint + is its final character. + + Falls back to the raw midpoint when no usable whitespace boundary exists, either + because `text` has none at all (a single giant token) or because the only + candidates were degenerate. That is still a correct, lossless split, just no + longer guaranteed word-safe for those cases. `text` must be at least two + characters, which `_split_bedrock_content` guarantees, so the midpoint itself + is never degenerate. + """ + midpoint: Final = len(text) // 2 + before: Final = max((found.end() for found in _BEDROCK_WHITESPACE.finditer(text, 0, midpoint)), default=None) + after_match: Final = _BEDROCK_WHITESPACE.search(text, midpoint) + candidates: Final = sorted( + (split for split in (before, after_match.end() if after_match else None) if split is not None), + key=lambda split: abs(split - midpoint), + ) + return next((split for split in candidates if 0 < split < len(text)), midpoint) + + @staticmethod + def _is_input_too_large_error(detail: object) -> bool: + """True if `detail` is an AWS error message for input exceeding the + per-request text-unit quota. + + Matched on the message rather than the status code on purpose: AWS is not + consistent about which error it raises for this. Observed against a live + guardrail with an active content-filter policy, an oversized request comes + back as a *ThrottlingException* (429) reading ``Input text size (3273 text + units) exceeds the maximum allowed (1000 text units) for the content filter + policy (Classic tier)``, while the documented failure mode is a + ValidationException (400). Keying off the message covers both. + + A guardrail *block* is also raised as an HTTPException with status 400, + but its ``detail`` is always a dict (built by + ``_get_http_exception_for_blocked_guardrail``); a non-200 API error's + ``detail`` is always the plain string returned by + ``_parse_bedrock_guardrail_error_response``. Checking ``isinstance(detail, + str)`` is therefore sufficient to never mistake a real block for a + too-large error. + """ + if not isinstance(detail, str): + return False + lowered: Final = detail.lower() + return any(substring in lowered for substring in _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS) + + @staticmethod + def _merge_bedrock_guardrail_responses( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> BedrockGuardrailResponse: + """Merge the per-chunk ApplyGuardrail responses of a chunked request into + one, so a caller cannot tell whether chunking happened. + + Only ever called with responses that all passed (a block raises + immediately from ``_apply_guardrail_content_with_chunking`` and is never + added to this list). ``action`` is only set on the merged response when + at least one chunk's raw response included it, and left absent otherwise + -- mirroring a real single-call response and matching what + ``_build_tracing_detail`` treats as "Bedrock didn't report an action". + + Fields this merge has no opinion on (``actionReason``, ``guardrailCoverage``, + ``blockedResponse``, anything AWS adds later) are carried over from the chunk + responses rather than dropped, so the response and the logged telemetry keep + the shape a single unchunked call returned. The merged keys below win. + + Per AWS's documented ApplyGuardrail contract, a single call's ``outputs`` + is positionally parallel to the ``content`` items *of that call*: an + entry per item when anything in the call was masked, or an empty list + when nothing in the whole call was masked. Downstream masking + (``_apply_masking_to_messages``) walks the merged ``outputs`` by a single + running index across the *original, unchunked* message list, so a later + chunk's masked text must land at the same global position it would have + if chunking had never happened. Naively concatenating each chunk's + ``outputs`` breaks that whenever a chunk had nothing masked (its empty + list would otherwise silently swallow its items' slots, shifting every + later chunk's masked text left onto the wrong message). So every + item -- masked or not -- always contributes exactly one entry here, + falling back to that item's own original (unmasked) text when its + chunk returned no output for it; a wholly-untouched result is then + collapsed back to an empty ``outputs`` list to match a real single-call + no-op response. A chunk that returns a nonzero output count not equal + to its item count is passed through as-is instead of guessed at, since + AWS's docs don't cover partial masking within one multi-item call. + """ + logical_units: Final = BedrockGuardrail._group_fragment_units(chunk_results) + per_unit_outputs: Final = tuple(BedrockGuardrail._merge_logical_unit_outputs(unit) for unit in logical_units) + merged_outputs: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list + output for outputs, _ in per_unit_outputs for output in outputs + ] + any_masked: Final = any(masked for _, masked in per_unit_outputs) + + actions: Final = tuple( + chunk_result.response.get("action") + for chunk_result in chunk_results + if isinstance(chunk_result.response.get("action"), str) + ) + merged_action: Final = ( + "GUARDRAIL_INTERVENED" if "GUARDRAIL_INTERVENED" in actions else (actions[-1] if actions else None) + ) + merged_assessments: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list + assessment + for chunk_result in chunk_results + for assessment in (chunk_result.response.get("assessments") or []) # mutable-ok: logged payload + ] + any_usage_reported: Final = any(chunk_result.response.get("usage") for chunk_result in chunk_results) + + merged: Final[BedrockGuardrailResponse] = cast( # cast-ok: TypedDict assembled from a comprehension + BedrockGuardrailResponse, + { # mutable-ok: builds the TypedDict payload + key: value for chunk_result in chunk_results for key, value in chunk_result.response.items() + }, + ) + if merged_action is not None: + merged["action"] = merged_action + if merged_outputs and any_masked: + merged["outputs"] = merged_outputs + merged["output"] = merged_outputs + if merged_assessments: + merged["assessments"] = merged_assessments + if any_usage_reported: + merged["usage"] = BedrockGuardrail._sum_bedrock_guardrail_usage(chunk_results) + return merged + + @staticmethod + def _sum_bedrock_guardrail_usage( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> BedrockGuardrailUsage: + """Sum each chunk's ``usage`` counters field-by-field into one totals dict. + + Keys are taken from the responses rather than from a fixed list, so a counter + this code does not know about (AWS has added several) is still summed and + reported instead of being silently dropped to zero.""" + chunk_usages: Final = tuple( + chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback + for chunk_result in chunk_results + ) + return cast( # cast-ok: TypedDict assembled from a comprehension + BedrockGuardrailUsage, + { # mutable-ok: builds the TypedDict payload + key: sum(usage.get(key) or 0 for usage in chunk_usages) + for key in dict.fromkeys(key for usage in chunk_usages for key in usage) + }, + ) + + @staticmethod + def _group_fragment_units( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> tuple[tuple[BedrockContentChunkResult, ...], ...]: + """Group consecutive text-fragment chunk results back into the one content + item each group came from, leaving every ordinary chunk result as a unit of + one. + + The group size is read off the results themselves rather than assumed, + because a single content item can be bisected repeatedly: two levels of + splitting yield four fragments for one item, not two. Assuming a fixed pair + here would emit two outputs for one message and shift every later message's + masked text onto the wrong message.""" + + def advance(carried: tuple[int, bool], result: BedrockContentChunkResult) -> tuple[int, bool]: + remaining, _ = carried + if remaining == 0: + return max(1, result.fragment_group_size) - 1, True + return remaining - 1, False + + starts: Final = tuple( + index + for index, (_, starts_unit) in enumerate(tuple(accumulate(chunk_results, advance, initial=(0, False)))[1:]) + if starts_unit + ) + return tuple(tuple(chunk_results[start:end]) for start, end in zip(starts, starts[1:] + (len(chunk_results),))) + + @staticmethod + def _merge_logical_unit_outputs( + unit: tuple[BedrockContentChunkResult, ...], + ) -> tuple[tuple[BedrockGuardrailOutput, ...], bool]: + """Reduce one logical unit (a fragment group of any size, or a single chunk + result) to the ``BedrockGuardrailOutput`` entries it contributes to the + merged response, plus whether any masking actually happened in it. + + Per AWS's documented ApplyGuardrail contract, a single call's + ``outputs`` is positionally parallel to the ``content`` items *of that + call*: an entry per item when anything in the call was masked, or an + empty list when nothing in the whole call was masked. Downstream + masking (``_apply_masking_to_messages``) walks the merged ``outputs`` + by a single running index across the *original, unchunked* message + list, so a later chunk's masked text must land at the same global + position it would have if chunking had never happened. So every item + -- masked or not -- always contributes exactly one entry here, falling + back to that item's own original (unmasked) text when its chunk + returned no output for it. A chunk that returns a nonzero output count + not equal to its item count is passed through as-is instead of guessed + at, since AWS's docs don't cover partial masking within one multi-item + call. + + A unit holding more than one result is a fragment group: every result in it + is one fragment of a single content item's text, so the group collapses to + one entry built from each fragment's masked text (or that fragment's own + original text where it came back unmasked), concatenated in order. This + holds for any group size, not only two. + """ + if len(unit) > 1: + + def fragment_outputs(result: BedrockContentChunkResult) -> tuple[BedrockGuardrailOutput, ...]: + return tuple(result.response.get("outputs") or result.response.get("output") or ()) + + def fragment_text(result: BedrockContentChunkResult) -> str: + source: Final = (result.content[0].get("text") or {}).get( # mutable-ok: read-only fallback + "text" + ) or "" + outputs: Final = fragment_outputs(result) + masked: Final = outputs[0].get("text") if outputs else None + return masked if masked is not None else source + + merged_text: Final = "".join(fragment_text(result) for result in unit) + any_masked: Final = any(fragment_outputs(result) for result in unit) + return (BedrockGuardrailOutput(text=merged_text),), any_masked + + (chunk_result,) = unit + chunk_outputs: Final = chunk_result.response.get("outputs") or chunk_result.response.get("output") or () + if len(chunk_outputs) == len(chunk_result.content): + return tuple(chunk_outputs), bool(chunk_outputs) + if not chunk_outputs: + return tuple( + BedrockGuardrailOutput( + text=(item.get("text") or {}).get("text") or "" # mutable-ok: read-only fallback + ) + for item in chunk_result.content + ), False + return tuple(chunk_outputs), True async def _sign_and_post( self, @@ -855,6 +1588,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, event_type: GuardrailEventHooks, start_time: "datetime", + log_transport_failure: bool = True, ) -> httpx.Response: """POST a signed Bedrock request, logging+raising on network/HTTP errors. @@ -862,6 +1596,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): transport-error handling cannot drift. Returns the raw ``httpx.Response`` on success (including non-2xx that httpx did not raise on); the 200-path logging, status and tracing stay with each caller because the two APIs report differently. + + ``log_transport_failure=False`` suppresses the ``guardrail_failed_to_respond`` + entry for a non-200 that is re-raised as an ``HTTPException``, for callers that + own consolidated per-request logging. The ApplyGuardrail path needs this: + ``AsyncHTTPHandler.post`` calls ``raise_for_status()``, so every non-200 lands + in this handler, and one logical request can legitimately produce several of + them (a too-large probe, then each rejected bisection level) while still + succeeding overall. Logging per attempt would report a recovered request as + several failures plus a success. + + The connection-level branch below (timeout, endpoint down) still logs + unconditionally: it re-raises the original exception rather than an + ``HTTPException``, so no consolidating caller catches it, and suppressing it + would drop the only record of the failure. """ try: return await self.async_handler.post( @@ -882,16 +1630,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): status_code, detail_message, ) = self._parse_bedrock_guardrail_error_response(err_response) - self.add_standard_logging_guardrail_information_to_request_data( - guardrail_provider=self.guardrail_provider, - guardrail_json_response={"error": detail_message}, - request_data=request_data or {}, - guardrail_status="guardrail_failed_to_respond", - start_time=start_time.timestamp(), - end_time=datetime.now(timezone.utc).timestamp(), - duration=(datetime.now(timezone.utc) - start_time).total_seconds(), - event_type=event_type, - ) + if log_transport_failure: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={ # mutable-ok: logging helper requires a dict + "error": detail_message + }, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + ) raise HTTPException(status_code=status_code, detail=detail_message) from e except HTTPException: raise @@ -900,7 +1651,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": str(e)}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1027,7 +1778,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": detail_message}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1043,7 +1794,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": str(e)}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1061,7 +1812,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response=self._sanitize_invoke_checks_response_for_logging(json_response), - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status=self._get_invoke_checks_status(bool(violations)), start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index 068a3ecf31b..facb822d00d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -19,7 +19,7 @@ request is sent with the ``X-Cisco-AI-Defense-API-Key`` header. import json import os -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -94,13 +94,13 @@ class _CiscoVerdict: is_safe: bool | None classifications: list[str] severity: str | None - rules: list[dict[str, Any]] + rules: list[dict[str, object]] explanation: str | None event_id: str | None action: str | None = None sanitized_text: str | None = None - sanitized_messages: list[dict[str, Any]] | None = None - sanitized_mcp_arguments: dict[str, Any] | None = None + sanitized_messages: list[dict[str, object]] | None = None + sanitized_mcp_arguments: dict[str, object] | None = None class CiscoAIDefenseGuardrailMissingSecrets(Exception): @@ -136,7 +136,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): api_base: str | None = None, inspection_type: str | None = None, inspect_path: str | None = None, - enabled_rules: list[dict[str, Any]] | None = None, + enabled_rules: Sequence[object] | None = None, integration_profile_id: str | None = None, integration_profile_version: str | None = None, integration_tenant_id: str | None = None, @@ -415,7 +415,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, - response: AsyncIterator[Any], + response: AsyncIterator[object], request_data: dict, ): """Buffer and inspect streaming chat output before delivery.""" @@ -437,7 +437,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self.guardrail_name, ) - all_chunks: Final[list[Any]] = [] + all_chunks: Final[list[object]] = [] try: async for chunk in response: all_chunks.append(chunk) @@ -497,7 +497,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): response_obj=assembled, ) except HTTPException as exc: - error_obj: dict[str, Any] = self._http_exception_to_error_obj(exc) + error_obj: dict[str, object] = self._http_exception_to_error_obj(exc) verbose_proxy_logger.warning( "Cisco AI Defense guardrail (%s): streaming response " "blocked — emitting SSE error event instead of " @@ -531,7 +531,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): for chunk in all_chunks: yield chunk - def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, Any]: + def _build_block_payload(self, context: _ScanContext, verdict: _CiscoVerdict) -> dict[str, object]: """Canonical block payload used across all four block paths. Same dict is the ``HTTPException.detail`` for chat / MCP request @@ -555,34 +555,34 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): "event_id": verdict.event_id, } - def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, Any]: + def _http_exception_to_error_obj(self, exc: HTTPException) -> dict[str, object]: """Wrap an ``HTTPException`` detail into the SSE ``error`` payload. For Cisco's own blocks the detail is already the canonical block payload, so this is a near-passthrough that just adds ``code`` / ``guardrail`` defaults for non-Cisco / unstructured details. """ - error_obj: dict[str, Any] = dict(exc.detail) if isinstance(exc.detail, dict) else {"message": str(exc.detail)} + error_obj: dict[str, object] = {**exc.detail} if isinstance(exc.detail, dict) else {"message": str(exc.detail)} error_obj.setdefault("message", error_obj.get("error", "Guardrail block")) error_obj.setdefault("code", exc.status_code) error_obj.setdefault("guardrail", self.guardrail_name) return error_obj @classmethod - def _streaming_content_was_modified(cls, original_chunks: list[Any], assembled: ModelResponse) -> bool: + def _streaming_content_was_modified(cls, original_chunks: Sequence[object], assembled: ModelResponse) -> bool: """Decide whether redact changed content or tool/function arguments.""" original_text: Final = cls._extract_streaming_chunk_scan_text(original_chunks) assembled_text: Final = " ".join(m.get("content", "") for m in cls._extract_response_messages(assembled)) return original_text != assembled_text @classmethod - def _extract_streaming_chunk_scan_text(cls, chunks: list[Any]) -> str: + def _extract_streaming_chunk_scan_text(cls, chunks: Sequence[object]) -> str: original_text = "" argument_text = "" for chunk in chunks: choices = getattr(chunk, "choices", None) or [] for c in choices: - delta = getattr(c, "delta", None) + delta: object | None = getattr(c, "delta", None) if delta is None: continue text = getattr(delta, "content", None) @@ -595,7 +595,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): args = cls._extract_tool_call_arguments(tc) if args: argument_text += args - fc = getattr(delta, "function_call", None) + fc: object | None = getattr(delta, "function_call", None) if fc is not None: args = cls._extract_function_call_arguments(fc) if args: @@ -673,7 +673,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): allow, WARNING for intervened/redacted, ERROR is left for upstream API failures. """ - fields: Final[dict[str, Any]] = { + fields: Final[dict[str, object]] = { "guardrail": self.guardrail_name, "surface": context.surface, "direction": context.direction, @@ -752,7 +752,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, direction: str = "input", response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: url: Final = f"{self.api_base}{self.inspect_path}" payload: Final = self._build_chat_payload(messages, request_data, user_api_key_dict) start_time: Final = datetime.now() @@ -784,7 +784,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): messages: list[dict[str, str]], request_data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: + ) -> dict[str, object]: return { "messages": messages, "metadata": self._build_metadata(request_data, user_api_key_dict), @@ -798,9 +798,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): async def _post_inspection( self, url: str, - payload: dict[str, Any], + payload: dict[str, object], surface: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: headers: Final = self._build_headers() verbose_proxy_logger.debug( "Cisco AI Defense guardrail: posting %s inspection to %s", @@ -856,8 +856,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - ) -> dict[str, Any]: - metadata: Final[dict[str, Any]] = {} + ) -> dict[str, object]: + metadata: Final[dict[str, object]] = {} user: Final = request_data.get("user") or getattr(user_api_key_dict, "user_id", None) if user: @@ -884,8 +884,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return metadata - def _build_config(self) -> dict[str, Any]: - config: Final[dict[str, Any]] = {} + def _build_config(self) -> dict[str, object]: + config: Final[dict[str, object]] = {} if self.enabled_rules: config["enabled_rules"] = self.enabled_rules if self.integration_profile_id: @@ -899,7 +899,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return config @staticmethod - def _normalize_rule(rule: object) -> dict[str, Any]: + def _normalize_rule(rule: object) -> dict[str, object]: """Coerce a user-supplied rule into the wire-shape dict Cisco expects. Accepts ``str``, ``dict``, and Pydantic model inputs. @@ -922,7 +922,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): rule = dumped if isinstance(rule, dict): - normalized: Final[dict[str, Any]] = {} + normalized: Final[dict[str, object]] = {} rule_name: Final = rule.get("rule_name") if rule_name: normalized["rule_name"] = rule_name @@ -950,7 +950,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): context: _ScanContext, start_time: datetime, response_obj: object = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Parse, log, and (optionally) raise/redact on the Cisco verdict. ``context.direction`` is ``"input"`` for request scans and ``"output"`` @@ -1119,10 +1119,10 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @classmethod def _sanitize_response_for_logging( cls, - inspect_response: dict[str, Any], + inspect_response: Mapping[str, object], surface: str, action: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """Drop bulky / privacy-sensitive fields, recursing into nested dicts. MCP verdicts are commonly nested under ``result``, so a @@ -1138,9 +1138,9 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return sanitized @classmethod - def _strip_sensitive_keys(cls, d: dict[str, Any]) -> dict[str, Any]: + def _strip_sensitive_keys(cls, d: Mapping[str, object]) -> dict[str, object]: """Recursively strip privacy-sensitive keys from a verdict dict.""" - out: Final[dict[str, Any]] = {} + out: Final[dict[str, object]] = {} for key, value in d.items(): if key.startswith("_") or key in cls._REDACTED_LOG_KEYS: continue @@ -1222,8 +1222,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_jsonrpc_error( - inspect_response: dict[str, Any], - ) -> dict[str, Any] | None: + inspect_response: Mapping[str, object], + ) -> dict[str, object] | None: """Detect a JSON-RPC error envelope inside an HTTP 200 response. The Cisco Inspect API can return ``{"error": {...}}`` (or nest one @@ -1270,7 +1270,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_sanitized_text( - inspect_response: dict[str, Any], + inspect_response: Mapping[str, object], ) -> str | None: """Pull ``sanitized_text`` (or camelCase variant) off the verdict.""" for key in ("sanitized_text", "sanitizedText"): @@ -1287,8 +1287,8 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _extract_sanitized_messages( - inspect_response: dict[str, Any], - ) -> list[dict[str, Any]] | None: + inspect_response: Mapping[str, object], + ) -> list[dict[str, object]] | None: """Pull a sanitized OpenAI-format messages array off the verdict. Cisco can return the rewrite under several keys; we accept any of @@ -1354,7 +1354,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_mcp_input( request_data: dict, sanitized_text: str | None, - sanitized_mcp_arguments: dict[str, Any] | None, + sanitized_mcp_arguments: dict[str, object] | None, ) -> bool: """Rewrite MCP request arguments in all locations the proxy reads.""" if sanitized_mcp_arguments is not None: @@ -1388,7 +1388,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, request_data: dict, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite chat request input (``messages`` or ``input``).""" if sanitized_messages and self._extract_tool_definition_text(request_data): @@ -1444,7 +1444,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): cls, request_data: dict, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: if sanitized_messages: instruction_text: Final = cls._instruction_text_from_messages(sanitized_messages) @@ -1457,7 +1457,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return False @classmethod - def _instruction_text_from_messages(cls, messages: list[dict[str, Any]]) -> str | None: + def _instruction_text_from_messages(cls, messages: list[dict[str, object]]) -> str | None: for message in messages: if not isinstance(message, dict): continue @@ -1468,7 +1468,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return None @classmethod - def _non_instruction_messages(cls, messages: list[dict[str, Any]] | None) -> list[dict[str, Any]] | None: + def _non_instruction_messages(cls, messages: list[dict[str, object]] | None) -> list[dict[str, object]] | None: if messages is None: return None return [ @@ -1499,7 +1499,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, response_obj: object, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite chat response (``ModelResponse`` or ``ResponsesAPIResponse``).""" if response_obj is None: @@ -1526,7 +1526,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_model_response_choices( choices: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Redact every returned choice, including tool-call/reasoning fields.""" if sanitized_messages: @@ -1570,7 +1570,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): def _redact_text_completion_choices( choices: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: """Rewrite ``/v1/completions`` text choices after Cisco redaction.""" replacement = sanitized_text @@ -1638,7 +1638,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): self, output_items: list, sanitized_text: str | None, - sanitized_messages: list[dict[str, Any]] | None, + sanitized_messages: list[dict[str, object]] | None, ) -> bool: replacement_text: str | None = sanitized_text if not replacement_text and sanitized_messages: @@ -1672,14 +1672,14 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): @staticmethod def _sanitized_messages_to_responses_input( - sanitized_messages: list[dict[str, Any]], - ) -> list[dict[str, Any]] | None: + sanitized_messages: list[dict[str, object]], + ) -> list[dict[str, object]] | None: """Convert chat-shape sanitized_messages to Responses API ``input``. Returns ``None`` if nothing usable could be converted, so the caller falls back to ``on_flagged_action``. """ - out: Final[list[dict[str, Any]]] = [] + out: Final[list[dict[str, object]]] = [] for m in sanitized_messages: if not isinstance(m, dict): continue @@ -1764,7 +1764,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): start_time: datetime | None = None, surface: str = "chat", direction: str = "input", - ) -> dict[str, Any]: + ) -> dict[str, object]: verbose_proxy_logger.error( "Cisco AI Defense guardrail (%s): API communication failed: %s", surface, @@ -2060,7 +2060,7 @@ class CiscoAIDefenseGuardrail(_CiscoAIDefenseMcpMixin, CustomGuardrail): return getattr(obj, key, None) @classmethod - def _field_list(cls, obj: object, key: str) -> list[Any]: + def _field_list(cls, obj: object, key: str) -> list[object]: value: Final = cls._field(obj, key) return value if isinstance(value, list) else [] diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 61220819d48..8bfd5cca58a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -148,6 +148,27 @@ def _restore_protected_messages( ] +def _build_compress_failure_detail(status_code: int, body: str) -> dict[str, object]: + """Build error details for failed /v1/compress responses. + + Adds troubleshooting hints for known deployment-related errors while + preserving the upstream status code and response body. + """ + if status_code == 404: + return { + "status_code": status_code, + "body": body, + "hint": ( + "The Headroom compression endpoint returned HTTP 404. " + "Verify that the configured Headroom endpoint is correct and that " + "the compression endpoint is available. If you are using a " + "self-hosted deployment, some deployments require enabling remote " + "compression (for example, HEADROOM_COMPRESS_ALLOW_REMOTE=1)." + ), + } + return {"status_code": status_code, "body": body} + + def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]: hashes: Final[list[str]] = [] for msg in messages: @@ -417,7 +438,7 @@ class HeadroomGuardrail(CustomGuardrail): self._handle_compress_failure( messages, "Headroom compression service returned an error", - {"status_code": e.response.status_code, "body": e.response.text}, + _build_compress_failure_detail(e.response.status_code, e.response.text), ), False, {}, @@ -449,7 +470,7 @@ class HeadroomGuardrail(CustomGuardrail): self._handle_compress_failure( messages, "Headroom compression service returned an error", - {"status_code": response.status_code, "body": response.text}, + _build_compress_failure_detail(response.status_code, response.text), ), False, {}, diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index c29da89b15f..5bbb01c6c8e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -8,8 +8,8 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint import copy import json -from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, Final +from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import HTTPException @@ -34,6 +34,9 @@ if TYPE_CHECKING: # Imported lazily at runtime (inside the streaming hook) to avoid a # module-level cyclic import with litellm.integrations.custom_guardrail. from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) # Call types that use NDJSON streaming (A2A); guardrail HTTPException is emitted as in-stream error A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) @@ -41,12 +44,35 @@ A2A_CALL_TYPES: Final = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME: Final = "unified_llm_guardrails" +class _EndpointTranslation(Protocol): + @property + def process_input_messages(self) -> "Callable[..., Awaitable[dict[str, object]]]": ... + + @property + def process_output_response(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def process_output_streaming_response(self) -> "Callable[..., Awaitable[object]]": ... + + @property + def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... + + +def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation: + return translation + + +def _chunk_choices(item: object) -> Sequence[object]: + choices: Final[Sequence[object]] = getattr(item, "choices", None) or [] + return choices + + class _StreamTerminated(Exception): """Internal signal that the incremental transform stream has already emitted its terminal chunks (block message or in-stream error) and must stop.""" -def _get_a2a_request_id(responses_so_far: list[Any], request_data: dict) -> str | None: +def _get_a2a_request_id(responses_so_far: Sequence[object], request_data: dict) -> str | None: """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting.""" for item in responses_so_far: if isinstance(item, dict) and "id" in item: @@ -138,7 +164,9 @@ class UnifiedLLMGuardrails(CustomLogger): except ValueError: return data # handle unmapped call types - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) _ensure_litellm_metadata(data, user_api_key_dict) @@ -156,7 +184,7 @@ class UnifiedLLMGuardrails(CustomLogger): async def async_moderation_hook( self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral - ) -> Any: + ) -> object: """ Runs in parallel to LLM API call Runs on only Input @@ -187,7 +215,9 @@ class UnifiedLLMGuardrails(CustomLogger): if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings: return data - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) _ensure_litellm_metadata(data, user_api_key_dict) @@ -202,7 +232,7 @@ class UnifiedLLMGuardrails(CustomLogger): data: dict, user_api_key_dict: UserAPIKeyAuth, response, - ) -> Any: + ) -> object: """ Runs on response from LLM API call @@ -271,7 +301,9 @@ class UnifiedLLMGuardrails(CustomLogger): ) return response - endpoint_translation: Final = endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation( + endpoint_guardrail_translation_mappings[CallTypes(call_type)]() + ) try: response = await endpoint_translation.process_output_response( @@ -299,10 +331,10 @@ class UnifiedLLMGuardrails(CustomLogger): async def _handle_streaming_block( self, exc: "ModifyResponseException", - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, stream_started: bool, - responses_so_far: list[Any], - ) -> AsyncGenerator[Any, None]: + responses_so_far: Sequence[object], + ) -> AsyncGenerator[object, None]: """ Terminate a streamed response cleanly when a guardrail blocks it. @@ -323,7 +355,7 @@ class UnifiedLLMGuardrails(CustomLogger): @staticmethod def _resolve_transform_call_type( user_api_key_dict: UserAPIKeyAuth, - mappings: dict, + mappings: Mapping[CallTypes, type["BaseTranslation"]], ) -> str | None: """Resolve the call type for the incremental_diff path, or None if the route is unresolvable / unsupported. @@ -356,9 +388,9 @@ class UnifiedLLMGuardrails(CustomLogger): self, exc: HTTPException, call_type: str | None, - responses_so_far: list[Any], + responses_so_far: Sequence[object], request_data: dict, - ) -> AsyncGenerator[Any, None]: + ) -> AsyncGenerator[object, None]: """Surface a mid-stream HTTPException. For A2A (NDJSON) call types the response has already started, so emit an in-stream JSON-RPC error chunk; otherwise re-raise so the proxy can report it. @@ -387,7 +419,7 @@ class UnifiedLLMGuardrails(CustomLogger): def _build_transform_chunk( self, *, - reference_chunk: Any, + reference_chunk: object, mutated_text_per_choice: dict[int, str], emitted_text_per_choice: dict[int, str], holdback_per_choice: dict[int, int], @@ -500,18 +532,18 @@ class UnifiedLLMGuardrails(CustomLogger): async def _emit_transform_round( self, *, - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, request_data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str, - reference_chunk: Any, - responses_so_far: list[Any], - responses_yielded: list[Any], + reference_chunk: object, + responses_so_far: Sequence[object], + responses_yielded: list[object], emitted_text_per_choice: dict[int, str], finish_reason_per_choice: dict[int, str | None], is_final: bool, - ) -> AsyncGenerator[Any, None]: + ) -> AsyncGenerator[object, None]: """Run one guardrail processing round and emit the resulting diff chunk. Raises ``_StreamTerminated`` (after emitting the terminal block message or @@ -564,14 +596,14 @@ class UnifiedLLMGuardrails(CustomLogger): self, *, guardrail_to_apply: CustomGuardrail, - response: Any, + response: AsyncIterable[object], request_data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str, sampling_rate: int, end_of_stream_only: bool, - mappings: dict, - ) -> AsyncGenerator[Any, None]: + mappings: Mapping[CallTypes, type["BaseTranslation"]], + ) -> AsyncGenerator[object, None]: """Emit guardrail text transformations as new deltas on the stream. Raw chunks are withheld and accumulated; on each sampled processing round @@ -580,15 +612,15 @@ class UnifiedLLMGuardrails(CustomLogger): synthetic chunk. A BLOCK terminates the stream via the shared block handler; an underflow surfaces as an HTTPException. """ - endpoint_translation: Final = mappings[CallTypes(call_type)]() - responses_so_far: Final[list[Any]] = [] - responses_yielded: Final[list[Any]] = [] + endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + responses_so_far: Final[list[object]] = [] + responses_yielded: Final[list[object]] = [] emitted_text_per_choice: Final[dict[int, str]] = {} finish_reason_per_choice: Final[dict[int, str | None]] = {} chunk_counter = 0 - last_chunk: Any | None = None + last_chunk: object | None = None - def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]: + def _round(reference_chunk: object, is_final: bool) -> AsyncGenerator[object, None]: return self._emit_transform_round( endpoint_translation=endpoint_translation, guardrail_to_apply=guardrail_to_apply, @@ -694,13 +726,13 @@ class UnifiedLLMGuardrails(CustomLogger): async def _inspect_full_response_for_block( self, *, - endpoint_translation: Any, + endpoint_translation: _EndpointTranslation, guardrail_to_apply: CustomGuardrail, request_data: dict, user_api_key_dict: UserAPIKeyAuth, - responses_so_far: list[Any], - responses_yielded: list[Any], - ) -> AsyncGenerator[Any, None]: + responses_so_far: Sequence[object], + responses_yielded: Sequence[object], + ) -> AsyncGenerator[object, None]: """Run the block-only guardrail inspection over the full assembled response (text + tool calls) so nothing bypasses the block decision. @@ -734,17 +766,17 @@ class UnifiedLLMGuardrails(CustomLogger): raise _StreamTerminated() @staticmethod - def _chunk_has_tool_calls(item: Any) -> bool: - for choice in getattr(item, "choices", None) or []: + def _chunk_has_tool_calls(item: object) -> bool: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) if getattr(delta, "tool_calls", None): return True return False @staticmethod - def _chunk_carries_text(item: Any) -> bool: + def _chunk_carries_text(item: object) -> bool: """True if any choice in this chunk has non-empty string ``delta.content``.""" - for choice in getattr(item, "choices", None) or []: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) content = getattr(delta, "content", None) if isinstance(content, str) and content != "": @@ -753,7 +785,7 @@ class UnifiedLLMGuardrails(CustomLogger): @staticmethod def _tool_call_passthrough_chunk( - item: Any, + item: object, finish_reason_per_choice: "dict[int, str | None] | None" = None, ) -> ModelResponseStream: """Copy of a chunk carrying tool calls with all text content stripped. @@ -772,7 +804,7 @@ class UnifiedLLMGuardrails(CustomLogger): redaction purpose. """ synthetic_choices: Final[list[StreamingChoices]] = [] - for choice in getattr(item, "choices", None) or []: + for choice in _chunk_choices(item): delta = getattr(choice, "delta", None) idx = getattr(choice, "index", 0) or 0 original_finish = getattr(choice, "finish_reason", None) @@ -801,15 +833,15 @@ class UnifiedLLMGuardrails(CustomLogger): ) @staticmethod - def _record_finish_reasons(item: Any, finish_reason_per_choice: dict[int, str | None]) -> None: - for choice in getattr(item, "choices", None) or []: + def _record_finish_reasons(item: object, finish_reason_per_choice: dict[int, str | None]) -> None: + for choice in _chunk_choices(item): finish_reason = getattr(choice, "finish_reason", None) if finish_reason is not None: finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason @staticmethod - def _chunk_has_finish_reason(item: Any) -> bool: - choices: Final = getattr(item, "choices", None) or [] + def _chunk_has_finish_reason(item: object) -> bool: + choices: Final = _chunk_choices(item) return any(getattr(choice, "finish_reason", None) is not None for choice in choices) async def async_post_call_streaming_iterator_hook( @@ -845,22 +877,22 @@ class UnifiedLLMGuardrails(CustomLogger): # Get streaming configuration. Resolution order (later wins): default # < guardrail attribute < guardrail_config dict < this callback's # optional_params. - def _streaming_flag(name: str, default: Any) -> Any: + def _streaming_flag(name: str, default: object) -> Any: value = default if guardrail_to_apply is not None: value = getattr(guardrail_to_apply, name, value) - config: Final = getattr(guardrail_to_apply, "guardrail_config", {}) + config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {}) if isinstance(config, dict): value = config.get(name, value) return self.optional_params.get(name, value) - sampling_rate: Final = _streaming_flag("streaming_sampling_rate", 5) + sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). - end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False) + end_of_stream_only: bool = _streaming_flag("streaming_end_of_stream_only", False) # "block_only" (default) drops guardrail text rewrites on the streaming # path; "incremental_diff" emits them as synthetic deltas (see # _run_incremental_transform_stream). - streaming_transform_mode: Final = _streaming_flag("streaming_transform_mode", "block_only") + streaming_transform_mode: Final[str] = _streaming_flag("streaming_transform_mode", "block_only") # Withhold every chunk until end-of-stream moderation passes, then # release the original chunks (clean) or only the block message # (blocked) -- moderating the whole response *before* any content @@ -868,7 +900,9 @@ class UnifiedLLMGuardrails(CustomLogger): # release the original chunks are replayed as-is, so a # content-rewriting guardrail (e.g. PII masking) would leak # unredacted content. Guarded below via mask_response_content. - buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default) + buffer_until_moderated: bool = _streaming_flag( + "streaming_buffer_until_moderated", buffer_until_moderated_default + ) if ( buffer_until_moderated @@ -939,9 +973,9 @@ class UnifiedLLMGuardrails(CustomLogger): # Infer call type from first chunk call_type = None chunk_counter = 0 - responses_so_far: Final[list[Any]] = [] - responses_yielded: Final[list[Any]] = [] - pending_end_of_stream_items: Final[list[Any]] = [] + responses_so_far: Final[list[object]] = [] + responses_yielded: Final[list[object]] = [] + pending_end_of_stream_items: Final[list[object]] = [] # Whether any real response chunk has been forwarded to the client. # Drives how a block terminates the stream: continue the in-progress # message (True) vs emit a standalone block message (False, buffered). diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py index 408260d8483..270c28b094e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/__init__.py @@ -18,6 +18,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" send_user_api_key_alias=litellm_params.send_user_api_key_alias, send_user_api_key_user_id=litellm_params.send_user_api_key_user_id, send_user_api_key_team_id=litellm_params.send_user_api_key_team_id, + timeout=litellm_params.timeout, guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, default_on=litellm_params.default_on, diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index c5c66988cb4..1aefa38ecf8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -22,9 +22,10 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.guardrails import LitellmParams from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel -GUARDRAIL_TIMEOUT: Final = 5 +DEFAULT_GUARDRAIL_TIMEOUT: Final = 5.0 class ZscalerAIGuard(CustomGuardrail): @@ -43,6 +44,7 @@ class ZscalerAIGuard(CustomGuardrail): send_user_api_key_alias: bool | None = None, send_user_api_key_user_id: bool | None = None, send_user_api_key_team_id: bool | None = None, + timeout: float | None = None, **kwargs, ): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) @@ -68,6 +70,7 @@ class ZscalerAIGuard(CustomGuardrail): if send_user_api_key_team_id is not None else os.getenv("SEND_USER_API_KEY_TEAM_ID", "False").lower() in ("true", "1") ) + self.timeout = self._resolve_timeout(timeout) verbose_proxy_logger.debug( "send_user_api_key_alias: %s, \n send_user_api_key_user_id:%s, \n send_user_api_key_team_id:%s", @@ -80,6 +83,29 @@ class ZscalerAIGuard(CustomGuardrail): verbose_proxy_logger.debug("ZscalerAIGuard Initializing ...") + @staticmethod + def _resolve_timeout(timeout: float | None) -> float: + """ + Resolve the effective per-request timeout, falling back to the default + when it is unset or non-positive. + """ + if timeout is None: + return DEFAULT_GUARDRAIL_TIMEOUT + + if timeout <= 0: + verbose_proxy_logger.warning( + "Ignoring non-positive Zscaler AI Guard timeout %s, using %s seconds", + timeout, + DEFAULT_GUARDRAIL_TIMEOUT, + ) + return DEFAULT_GUARDRAIL_TIMEOUT + + return timeout + + def update_in_memory_litellm_params(self, litellm_params: "LitellmParams") -> None: + super().update_in_memory_litellm_params(litellm_params) + self.timeout = self._resolve_timeout(litellm_params.timeout) + @staticmethod def _resolve_metadata_value(request_data: dict | None, key: str) -> str | None: """ @@ -267,7 +293,7 @@ class ZscalerAIGuard(CustomGuardrail): f"{url}", headers=headers, json=data, - timeout=GUARDRAIL_TIMEOUT, + timeout=self.timeout, ) response.raise_for_status() return response diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 190d19f3d52..0d23e19f88d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -20,6 +20,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): content_filter_threshold=litellm_params.content_filter_threshold, prompt_attack_threshold=litellm_params.prompt_attack_threshold, pii_confidence_threshold=litellm_params.pii_confidence_threshold, + chunk_budget_chars=litellm_params.chunk_budget_chars, default_on=litellm_params.default_on, disable_exception_on_block=litellm_params.disable_exception_on_block, mask_request_content=litellm_params.mask_request_content, diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index dd61cad15a1..a64ed764a67 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -26,7 +26,8 @@ Usage: import base64 import json -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -36,7 +37,10 @@ from litellm.llms.litellm_proxy.skills.prompt_injection import ( SkillPromptInjectionHandler, ) from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth -from litellm.types.utils import CallTypes, CallTypesLiteral +from litellm.types.utils import CallTypes, CallTypesLiteral, LLMResponseTypes + +if TYPE_CHECKING: + from litellm.llms.litellm_proxy.skills.sandbox_executor import SkillsSandboxExecutor class SkillsInjectionHook(CustomLogger): @@ -99,7 +103,7 @@ class SkillsInjectionHook(CustomLogger): verbose_proxy_logger.debug("SkillsInjectionHook: Processing %s skills", len(skills)) litellm_skills: Final[list[LiteLLM_SkillsTable]] = [] - anthropic_skills: Final[list[dict[str, Any]]] = [] + anthropic_skills: Final[list[dict[str, object]]] = [] # Separate skills by prefix for skill in skills: @@ -324,9 +328,9 @@ class SkillsInjectionHook(CustomLogger): async def async_post_call_success_deployment_hook( self, request_data: dict, - response: Any, + response: LLMResponseTypes, call_type: CallTypes | None, - ) -> Any | None: + ) -> LLMResponseTypes | None: """ Post-call hook to handle automatic code execution. @@ -372,7 +376,7 @@ class SkillsInjectionHook(CustomLogger): # Check if any tool call needs execution (litellm_code_execution or skill tool) has_executable_tool = False for tc in tool_calls: - tool_name = tc.get("name", "") + tool_name: str = tc.get("name", "") # Execute if it's litellm_code_execution OR a skill tool (litellm_skill_xxx) if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value or tool_name.startswith(LITELLM_SKILL_ID_PREFIX): has_executable_tool = True @@ -441,7 +445,7 @@ class SkillsInjectionHook(CustomLogger): data: dict, response: Any, skill_files: dict[str, bytes], - ) -> Any: + ) -> LLMResponseTypes | None: """ Execute the code execution loop for messages API (Anthropic format). @@ -466,7 +470,7 @@ class SkillsInjectionHook(CustomLogger): max_tokens: Final = data.get("max_tokens", 4096) executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[dict[str, object]]] = [] current_response = response for iteration in range(self.max_iterations): @@ -511,9 +515,9 @@ class SkillsInjectionHook(CustomLogger): # Process tool calls tool_results = [] for tc in tool_calls: - tool_name = tc.get("name", "") + tool_name: str = tc.get("name", "") tool_id = tc.get("id", "") - tool_input = tc.get("input", {}) + tool_input: Mapping[str, str] = tc.get("input", {}) # Execute if it's litellm_code_execution OR a skill tool if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value: @@ -561,8 +565,8 @@ class SkillsInjectionHook(CustomLogger): self, code: str, skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute code in sandbox and return result string.""" try: @@ -574,7 +578,8 @@ class SkillsInjectionHook(CustomLogger): # Collect generated files if exec_result.get("files"): - for f in exec_result["files"]: + files: Final[Sequence[Mapping[str, str]]] = exec_result["files"] + for f in files: generated_files.append( { "name": f["name"], @@ -595,10 +600,10 @@ class SkillsInjectionHook(CustomLogger): async def _execute_skill_tool( self, tool_name: str, - tool_input: dict[str, Any], + tool_input: Mapping[str, str], skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute a skill tool by generating and running code based on skill content.""" # Generate code based on available skill modules @@ -670,7 +675,7 @@ print('No executable skill module found') data: dict, response: Any, skill_files: dict[str, bytes], - ) -> Any: + ) -> LLMResponseTypes: """ Execute the code execution loop until model gives final response. @@ -704,7 +709,7 @@ print('No executable skill module found') kwargs: Final = {k: v for k, v in data.items() if k not in _EXCLUDED_ACOMPLETION_KEYS} executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout) - generated_files: Final[list[dict[str, Any]]] = [] + generated_files: Final[list[dict[str, object]]] = [] current_response: Any = response for iteration in range(self.max_iterations): @@ -713,7 +718,7 @@ print('No executable skill module found') stop_reason = current_response.choices[0].finish_reason # Build assistant message for conversation history - assistant_msg_dict: dict[str, Any] = { + assistant_msg_dict: dict[str, object] = { "role": "assistant", "content": assistant_message.content, } @@ -781,13 +786,13 @@ print('No executable skill module found') self, tool_call: Any, skill_files: dict[str, bytes], - executor: Any, - generated_files: list[dict[str, Any]], + executor: "SkillsSandboxExecutor", + generated_files: list[dict[str, object]], ) -> str: """Execute a litellm_code_execution tool call and return result string.""" try: args: Final = json.loads(tool_call.function.arguments) - code: Final = args.get("code", "") + code: Final[str] = args.get("code", "") verbose_proxy_logger.debug("SkillsInjectionHook: Executing code (%s chars)", len(code)) @@ -802,7 +807,8 @@ print('No executable skill module found') # Collect generated files if exec_result.get("files"): tool_result += "\n\nGenerated files:" - for f in exec_result["files"]: + files: Final[Sequence[Mapping[str, str]]] = exec_result["files"] + for f in files: file_content = base64.b64decode(f["content_base64"]) generated_files.append( { @@ -830,8 +836,8 @@ print('No executable skill module found') def _attach_files_to_response( self, response: Any, - generated_files: list[dict[str, Any]], - ) -> Any: + generated_files: list[dict[str, object]], + ) -> LLMResponseTypes: """ Attach generated files to the response object. @@ -841,11 +847,13 @@ print('No executable skill module found') if not generated_files: return response + raw_response: Final = response + # Handle dict response (Anthropic/messages API format) if isinstance(response, dict): response["_litellm_generated_files"] = generated_files verbose_proxy_logger.debug("SkillsInjectionHook: Attached %s files to dict response", len(generated_files)) - return response + return raw_response # Handle object response (OpenAI format) try: diff --git a/litellm/proxy/hooks/max_budget_limiter.py b/litellm/proxy/hooks/max_budget_limiter.py index b8a69705b99..eaf37b0bcf1 100644 --- a/litellm/proxy/hooks/max_budget_limiter.py +++ b/litellm/proxy/hooks/max_budget_limiter.py @@ -32,9 +32,12 @@ class _PROXY_MaxBudgetLimiter(CustomLogger): if max_budget is None or user_id is None: return - # Personal budget applies only to non-team requests, matching - # the explicit team-key exemption in common_checks section 4.1. - if user_api_key_dict.team_id is not None: + from litellm.proxy.proxy_server import general_settings + + if ( + user_api_key_dict.team_id is not None + and general_settings.get("apply_user_budget_to_team_keys") is not True + ): return # The reservation path admits at the strict-`<` boundary and diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b00ba35b14e..83ae59ef050 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -151,6 +151,20 @@ LITELLM_METADATA_ROUTES: Final = ( "files", ) +LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( + { + "mask_input", + "mask_output", + "session_id", + "trace_id", + "trace_metadata", + "trace_name", + "trace_release", + "trace_user_id", + "trace_version", + } +) + _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "proxy_server_request", "standard_logging_object", @@ -458,6 +472,18 @@ def _get_metadata_variable_name(request: Request) -> str: return "metadata" +def _promoted_trace_control_fields( + requester_metadata: Mapping[str, Any], + litellm_metadata: Mapping[str, Any], +) -> tuple[tuple[str, Any], ...]: + """Return the caller's trace-control fields that ``litellm_metadata`` does not already set.""" + return tuple( + (key, value) + for key, value in requester_metadata.items() + if key in LITELLM_TRACE_CONTROL_METADATA_FIELDS and key not in litellm_metadata + ) + + def _extract_generic_session_id_from_headers( normalized: dict[str, str], ) -> str | None: @@ -1670,6 +1696,13 @@ async def add_litellm_data_to_request( # paths may read from it. if "metadata" in data and isinstance(data["metadata"], dict): data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"]) + if _metadata_variable_name == "litellm_metadata": + data[_metadata_variable_name].update( + _promoted_trace_control_fields( + requester_metadata=data[_metadata_variable_name]["requester_metadata"], + litellm_metadata=data[_metadata_variable_name], + ) + ) # Merge litellm_metadata into the metadata variable (preserving existing # values). Runs after the user_api_key_* / _pipeline_managed_guardrails diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 141094f4d4c..d4221845b0c 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -4,8 +4,9 @@ AUTO ROUTER MANAGEMENT ENDPOINTS POST /auto_router/test_routing - Route one prompt through an unsaved complexity-router config """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final from pydantic import BaseModel, TypeAdapter @@ -260,6 +261,7 @@ async def preview_auto_router_routing( class _SessionAggRow(BaseModel): router_name: str router_type: str + tier_turns: Mapping[str, int] sessions: int turns: int unordered_turns: int @@ -284,6 +286,23 @@ class _SessionAggRow(BaseModel): _SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow]) _BENCHMARKS_SQL: Final = """ +WITH windowed AS ( + SELECT * FROM "LiteLLM_AutoRouterSession" + WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp +), +tier_maps AS ( + SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns + FROM ( + SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns + FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv + GROUP BY router_name, router_type, kv.key + ) per_tier + GROUP BY router_name, router_type +) +SELECT + agg.*, + COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns +FROM ( SELECT router_name, router_type, @@ -306,10 +325,11 @@ SELECT COALESCE(SUM(spend), 0)::float8 AS spend, COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds -FROM "LiteLLM_AutoRouterSession" -WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp +FROM windowed GROUP BY router_name, router_type -ORDER BY SUM(spend) DESC +) agg +LEFT JOIN tier_maps USING (router_name, router_type) +ORDER BY agg.spend DESC """ @@ -366,6 +386,7 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow: return _SessionAggRow( router_name="", router_type="", + tier_turns=MappingProxyType({}), sessions=sum(row.sessions for row in rows), turns=sum(row.turns for row in rows), unordered_turns=sum(row.unordered_turns for row in rows), @@ -443,6 +464,7 @@ async def get_auto_router_benchmarks( AutoRouterBenchmarkGroup( router_name=row.router_name, router_type=row.router_type, + tier_turns=row.tier_turns, **_benchmark_totals(row).model_dump(), ) for row in rows diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 640a735c916..abc5d3e53ff 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,7 +17,7 @@ import json import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, cast +from typing import Any, Final, Literal, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a5078c50fc0..2de1d177b33 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,7 +18,7 @@ import os import re import secrets import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast @@ -171,8 +171,12 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ... + async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ... + async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ... + async def update( self, *, @@ -181,6 +185,10 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): ) -> _PrismaRowT | None: ... +class _TxTables(Protocol): + litellm_proxymodeltable: _PrismaTableActions[object] + + def _prisma_table( repository: BaseRepository[_RepositoryModelT], ) -> _PrismaTableActions[_RepositoryModelT]: @@ -1650,9 +1658,12 @@ async def generate_key_fn( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - if user_custom_key_generate is not None: - if inspect.iscoroutinefunction(user_custom_key_generate): - result: Final = await user_custom_key_generate(data) + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( + user_custom_key_generate + ) + if custom_key_generate_hook is not None: + if inspect.iscoroutinefunction(custom_key_generate_hook): + result: Final = await custom_key_generate_hook(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision: Final = result.get("decision", True) @@ -1847,9 +1858,10 @@ async def generate_service_account_key_fn( verbose_proxy_logger.debug("entered /key/generate") - if user_custom_key_generate is not None: - if inspect.iscoroutinefunction(user_custom_key_generate): - result: Final = await user_custom_key_generate(data) + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate + if custom_key_generate_hook is not None: + if inspect.iscoroutinefunction(custom_key_generate_hook): + result: Final = await custom_key_generate_hook(data) else: raise ValueError("user_custom_key_generate must be a coroutine") decision: Final = result.get("decision", True) @@ -1918,7 +1930,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) casted_metadata[reserved_field] = existing_value - data_json: Final = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True) try: for k, v in data_json.items(): @@ -2179,7 +2191,7 @@ async def _process_single_key_update( llm_router: Router | None, user_custom_key_update: Callable | None = None, existing_key_row: LiteLLM_VerificationToken | None = None, -) -> dict[str, Any]: +) -> dict[str, object]: """ Process a single key update with all validations and checks. @@ -2722,9 +2734,10 @@ async def update_key_fn( ) # Custom key update hook - if user_custom_key_update is not None: - if inspect.iscoroutinefunction(user_custom_key_update): - result: Final = await user_custom_key_update(data) + custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update + if custom_key_update_hook is not None: + if inspect.iscoroutinefunction(custom_key_update_hook): + result: Final = await custom_key_update_hook(data) else: raise ValueError("user_custom_key_update must be a coroutine") decision: Final = result.get("decision", True) @@ -3528,11 +3541,34 @@ async def info_key_fn( ): """ Retrieve information about a key. + Parameters: - key: Optional[str] = Query parameter representing the key in the request - user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key + - key: str | None (query parameter) - The key to look up. Accepts the plaintext key or its hash. + Defaults to the key in the Authorization header. + Returns: - Dict containing the key and its associated information + - key: str - The key that was looked up, echoed back as it was passed in + - info: dict - The key's row, minus the hashed token + - key_alias: str | None - User-friendly key alias + - spend: float - Amount spent by the key. When budget_duration is set this covers only the + current budget window, not the key's lifetime + - max_budget: float | None - Max budget for the key, enforced against spend + - budget_duration: str | None - Budget reset period ("30d", "1h", etc.) + - budget_reset_at: datetime | None - When the current budget window ends and spend is next + reset to 0, not when it was last reset. Reset times snap to standard boundaries in the + configured timezone (30d and 1mo land on the 1st of the month, 7d on Monday, 1h on the + hour), so subtracting budget_duration from it does not give the window's start + - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + - model_max_budget_usage: dict | None - Current-window spend per model, present only when + the key has per-model budgets + - models: list - Model_name's the key is allowed to call + - tpm_limit / rpm_limit: int | None - Tokens and requests per minute limits + - metadata: dict - Metadata for the key, e.g. {"team": "core-infra"} + - blocked: bool | None - Whether the key is blocked + - expires: datetime | None - When the key stops authenticating requests + - last_active: datetime | None - When the key was last used + - object_permission: dict | None - Resolved vector store / MCP permissions when the key has + an object_permission_id Example Curl: ``` @@ -4066,10 +4102,11 @@ async def delete_verification_tokens( failed_tokens: list = [] try: if prisma_client: - tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many(where={"token": {"in": tokens}}) + hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens] + tokens = hashed_tokens + _keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await _prisma_table( + VerificationTokenRepository(prisma_client) + ).find_many(where={"token": {"in": hashed_tokens}}) if len(_keys_being_deleted) == 0: raise HTTPException( @@ -4268,7 +4305,7 @@ async def _rotate_master_key( if models: decrypted_models: Final = proxy_config.decrypt_model_list_from_db(new_models=models) verbose_proxy_logger.debug("ABLE TO DECRYPT MODELS - len(decrypted_models): %s", len(decrypted_models)) - new_models: Final = [] + new_models: Final[list[dict[str, object]]] = [] for model in decrypted_models: new_model = await _add_model_to_db( model_params=Deployment(**model), @@ -4283,7 +4320,8 @@ async def _rotate_master_key( _dumped["model_info"] = prisma.Json(_dumped["model_info"]) new_models.append(_dumped) verbose_proxy_logger.debug("Resetting proxy model table") - async with prisma_client.db.tx() as tx: + async with prisma_client.db.tx() as tx_ctx: + tx: Final[_TxTables] = tx_ctx await tx.litellm_proxymodeltable.delete_many() verbose_proxy_logger.debug("Creating %s models", len(new_models)) await tx.litellm_proxymodeltable.create_many( @@ -4607,7 +4645,7 @@ async def _execute_virtual_key_regeneration( _validate_key_alias_format(key_alias=new_key_alias) verbose_proxy_logger.debug("non_default_values: %s", non_default_values) update_data.update(non_default_values) - update_data = prisma_client.jsonify_object(data=update_data) + jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data) # If grace period set, insert deprecated key so old key remains valid await _insert_deprecated_key( @@ -4619,9 +4657,9 @@ async def _execute_virtual_key_regeneration( updated_token: Final = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, - data=update_data, + data=jsonified_update_data, ) - updated_token_dict: Final = dict(updated_token) if updated_token is not None else {} + updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {} updated_token_dict["key"] = new_token updated_token_dict["token_id"] = updated_token_dict.pop("token") @@ -5566,7 +5604,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -5579,7 +5617,7 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 @@ -5672,7 +5710,7 @@ def _build_key_filter_conditions( agent_id: str | None = None, use_substring_matching: bool = False, expires_filter: str | None = None, -) -> dict[str, str | dict[str, Any] | list[dict[str, Any]]]: +) -> Mapping[str, object]: """Build filter conditions for key listing. Visibility rules: @@ -5684,14 +5722,14 @@ def _build_key_filter_conditions( so former members cannot see service accounts they created after leaving. """ # Prepare filter conditions - where: dict[str, str | dict[str, Any] | list[dict[str, Any]]] = {} + where: dict[str, object] = {} where.update(_get_condition_to_filter_out_ui_session_tokens()) # Build the OR conditions for user's keys and admin team keys - or_conditions: Final[list[dict[str, Any]]] = [] + or_conditions: Final[list[dict[str, object]]] = [] # Base conditions for user's own keys - user_condition: Final[dict[str, Any]] = {} + user_condition: Final[dict[str, object]] = {} if user_id and isinstance(user_id, str): if use_substring_matching: user_condition["user_id"] = { @@ -5761,7 +5799,7 @@ def _build_key_filter_conditions( # Apply team_id, project_id and access_group_id as global AND filters so they # narrow results across all visibility conditions (own keys, team keys, etc.) - global_filters: tuple[dict[str, Any], ...] = ( + global_filters: Final[tuple[dict[str, object], ...]] = ( *( ( {"key_alias": {"contains": key_alias, "mode": "insensitive"}} @@ -5782,7 +5820,7 @@ def _build_key_filter_conditions( else () ), ) - combined_where = {"AND": [where, *global_filters]} if global_filters else where + combined_where: Final[Mapping[str, object]] = {"AND": [where, *global_filters]} if global_filters else where verbose_proxy_logger.debug("Filter conditions: %s", combined_where) return combined_where @@ -5963,7 +6001,7 @@ async def _list_key_helper( ) -def _get_condition_to_filter_out_ui_session_tokens() -> dict[str, Any]: +def _get_condition_to_filter_out_ui_session_tokens() -> Mapping[str, object]: """ Condition to filter out UI session tokens """ @@ -6372,7 +6410,7 @@ async def _can_user_query_key_info( async def test_key_logging( user_api_key_dict: UserAPIKeyAuth, request: Request, - key_logging: list[dict[str, Any]], + key_logging: Sequence[Mapping[str, str]], ) -> LoggingCallbackStatus: """ Test the key-based logging diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 71407c89813..c2087005863 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,9 +13,9 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import datetime import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from json import JSONDecodeError -from typing import Any, Final, Literal, cast +from typing import Final, Literal, Protocol, cast from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, ValidationError @@ -78,9 +78,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, - DeploymentTypedDict, GenericLiteLLMParams, - LiteLLMParamsTypedDict, updateDeployment, ) from litellm.utils import get_utc_datetime @@ -104,10 +102,80 @@ class UpdatePublicModelGroupsRequest(BaseModel): model_config = ConfigDict(extra="forbid") +class _ProxyModelRow(Protocol): + model_id: str + model_name: str + model_info: Mapping[str, object] | None + + def model_dump_json(self, *, exclude_none: bool = False) -> str: ... + + +class _ProxyModelTable(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ... + + def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ... + + def delete(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ... + + def delete_many(self, *, where: Mapping[str, object]) -> Awaitable[int]: ... + + +class _TxModelTables(Protocol): + litellm_proxymodeltable: _ProxyModelTable + + +class _TeamRow(Protocol): + models: Sequence[str] + + def model_dump(self) -> Mapping[str, object]: ... + + +class _TeamTable(Protocol): + def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ... + + def update( + self, *, where: Mapping[str, object], data: Mapping[str, object], include: Mapping[str, bool] + ) -> Awaitable[LiteLLM_TeamTable]: ... + + +class _TeamIdRef(Protocol): + team_id: str + + +class _ModelAliasRow(Protocol): + id: int + model_aliases: dict[str, str] + team: _TeamIdRef | None + + +class _ModelAliasTable(Protocol): + def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRow]]: ... + + def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ... + + +def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable: + return ModelRepository(prisma_client).table + + +def _repo_team_table(prisma_client: PrismaClient) -> _TeamTable: + return TeamRepository(prisma_client).table + + +def _db_team_table(prisma_client: PrismaClient) -> _TeamTable: + return prisma_client.db.litellm_teamtable + + +def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable: + return ModelTableRepository(prisma_client).table + + async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None: db_model: Final = cast( BaseModel | None, - await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id}), + await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id}), ) if not db_model: @@ -166,14 +234,9 @@ def _raise_on_strategy_router_write_violation( def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: - merged_deployment_dict: Final = DeploymentTypedDict( - model_name=db_model.model_name, - litellm_params=LiteLLMParamsTypedDict(**db_model.litellm_params.model_dump(exclude_none=True)), - model_info=db_model.model_info.model_dump(exclude_none=True), - ) - # update model name - if updated_patch.model_name: - merged_deployment_dict["model_name"] = updated_patch.model_name + 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) # update litellm params if updated_patch.litellm_params: @@ -182,13 +245,11 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr k: encrypt_value_helper(v) for k, v in updated_patch.litellm_params.model_dump(exclude_none=True).items() } - merged_deployment_dict["litellm_params"].update(encrypted_params) + merged_litellm_params.update(encrypted_params) # update model info if updated_patch.model_info: - if "model_info" not in merged_deployment_dict: - merged_deployment_dict["model_info"] = {} - merged_deployment_dict["model_info"].update(updated_patch.model_info.model_dump(exclude_none=True)) + merged_model_info.update(updated_patch.model_info.model_dump(exclude_none=True)) # Honor explicit-null clears LAST, after both merges, so a model_info blob the UI # passes through (which today re-sends the OLD pricing on every save) cannot @@ -202,29 +263,25 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr if updated_patch.litellm_params: for field in updated_patch.litellm_params.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.litellm_params, field) is None: - merged_deployment_dict["litellm_params"].pop(field, None) - merged_deployment_dict.get("model_info", {}).pop(field, None) + merged_litellm_params.pop(field, None) + merged_model_info.pop(field, None) if updated_patch.model_info: for field in updated_patch.model_info.model_fields_set: if field in SPECIAL_MODEL_INFO_PARAMS and getattr(updated_patch.model_info, field) is None: - merged_deployment_dict["model_info"].pop(field, None) - merged_deployment_dict.get("litellm_params", {}).pop(field, None) + merged_model_info.pop(field, None) + merged_litellm_params.pop(field, None) # convert to prisma compatible format - prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel() - if "model_name" in merged_deployment_dict: - prisma_compatible_model_dict["model_name"] = merged_deployment_dict["model_name"] + for key, value in merged_model_info.items(): + if isinstance(value, datetime.datetime): + merged_model_info[key] = value.isoformat() - if "litellm_params" in merged_deployment_dict: - prisma_compatible_model_dict["litellm_params"] = json.dumps(merged_deployment_dict["litellm_params"]) - - if "model_info" in merged_deployment_dict: - model_info: Final = merged_deployment_dict["model_info"] - for key, value in model_info.items(): - if isinstance(value, datetime.datetime): - model_info[key] = value.isoformat() - prisma_compatible_model_dict["model_info"] = json.dumps(model_info) + prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( + model_name=merged_model_name, + litellm_params=json.dumps(merged_litellm_params), + model_info=json.dumps(merged_model_info), + ) if updated_patch.blocked is not None: prisma_compatible_model_dict["blocked"] = updated_patch.blocked @@ -338,7 +395,7 @@ async def patch_model( update_data["updated_at"] = cast(str, get_utc_datetime()) # Perform partial update - updated_model: Final = await ModelRepository(prisma_client).table.update( + updated_model: Final = await _proxy_model_table(prisma_client).update( where={"model_id": model_id}, data=update_data, ) @@ -769,8 +826,8 @@ async def _setup_new_team_model_assignment( async def _get_team_deployments( - team_id: str, prisma_client: PrismaClient, table: Any | None = None -) -> list[LiteLLM_ProxyModelTable]: + team_id: str, prisma_client: PrismaClient, table: _ProxyModelTable | None = None +) -> Sequence[_ProxyModelRow]: """ Fetch all deployments for a given team_id from the database. @@ -785,7 +842,7 @@ async def _get_team_deployments( existing transaction. """ prefix: Final = f"model_name_{team_id}_" - table = table or ModelRepository(prisma_client).table + table = table or _proxy_model_table(prisma_client) response: Final = await table.find_many( where={ "model_name": {"startswith": prefix}, @@ -806,7 +863,7 @@ async def _get_team_deployments( async def delete_team_models( team_ids: list[str], prisma_client: PrismaClient, - llm_router: Any | None, + llm_router: Router | None, ) -> list[str]: """ Delete every BYOK model owned by the given teams, from the DB and the router. @@ -820,7 +877,8 @@ async def delete_team_models( Returns the model_ids that were deleted. """ deleted_model_ids: Final[list[str]] = [] - async with prisma_client.db.tx() as tx: + async with prisma_client.db.tx() as tx_ctx: + tx: Final[_TxModelTables] = tx_ctx for team_id in team_ids: rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable) model_ids = [row.model_id for row in rows] @@ -920,11 +978,11 @@ async def _remove_unbacked_team_models( if not names_to_remove: return - existing_team_row: Final = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id}) + existing_team_row: Final = await _db_team_table(prisma_client).find_unique(where={"team_id": team_id}) if existing_team_row is None: return - updated_team_row: Final[LiteLLM_TeamTable] = await prisma_client.db.litellm_teamtable.update( + updated_team_row: Final[LiteLLM_TeamTable] = await _db_team_table(prisma_client).update( where={"team_id": team_id}, data={"models": [model for model in existing_team_row.models if model not in names_to_remove]}, include={"object_permission": True}, @@ -953,7 +1011,7 @@ async def _update_existing_team_model_assignment( """ def _get_team_public_model_name( - model_info: dict | str | None, + model_info: object, ) -> str | None: parsed: Final = model_info_as_mapping(model_info) if parsed is None: @@ -1062,7 +1120,7 @@ class ModelManagementAuthChecks: detail={"error": CommonProxyErrors.not_premium_user.value}, ) - _existing_team_row: Final = await TeamRepository(prisma_client).table.find_unique( + _existing_team_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) @@ -1091,7 +1149,7 @@ class ModelManagementAuthChecks: ) -> Literal[True]: ## Check team model auth if model_params.model_info is not None and model_params.model_info.team_id is not None: - team_obj_row: Final = await TeamRepository(prisma_client).table.find_unique( + team_obj_row: Final = await _repo_team_table(prisma_client).find_unique( where={"team_id": model_params.model_info.team_id} ) if team_obj_row is None: @@ -1192,7 +1250,7 @@ async def delete_model( - store keys separately """ # encrypt litellm params # - result: Final = await ModelRepository(prisma_client).table.delete(where={"model_id": model_info.id}) + result: Final = await _proxy_model_table(prisma_client).delete(where={"model_id": model_info.id}) if result is None: raise HTTPException( @@ -1265,9 +1323,9 @@ async def delete_team_model_alias( Returns: - List of team id + model alias pairs that were removed """ - team_model_aliases: Final = await ModelTableRepository(prisma_client).table.find_many(include={"team": True}) + team_model_aliases: Final = await _model_alias_table(prisma_client).find_many(include={"team": True}) tasks: Final = [] - removed_model_aliases: Final = [] + removed_model_aliases: Final[list[tuple[str, str]]] = [] for team_model_alias in team_model_aliases: model_aliases = team_model_alias.model_aliases # {"alias": "public model name"} id = team_model_alias.id @@ -1278,7 +1336,7 @@ async def delete_team_model_alias( removed_model_aliases.append((team_model_alias.team.team_id, key)) del model_aliases[key] tasks.append( - ModelTableRepository(prisma_client).table.update( + _model_alias_table(prisma_client).update( where={"id": id}, data={"model_aliases": json.dumps(model_aliases)}, ) @@ -1492,7 +1550,7 @@ async def update_model( }, ) - _model_id = None + _model_id: str | None = None _model_info: Final = getattr(model_params, "model_info", None) if _model_info is None: raise Exception("model_info not provided") @@ -1551,11 +1609,11 @@ async def update_model( else: pass - _data: Final[dict] = { + _data: Final[dict[str, str]] = { "litellm_params": json.dumps(merged_dictionary), "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, } - model_response: Final = await ModelRepository(prisma_client).table.update( + model_response: Final = await _proxy_model_table(prisma_client).update( where={"model_id": _model_id}, data=_data, ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fe5a0e06d2e..97f494c51de 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -15,11 +15,11 @@ import math import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Annotated, Final, Protocol, TypeVar, cast +from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import litellm from litellm._logging import verbose_proxy_logger @@ -29,6 +29,7 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( UI_TEAM_ID, BlockTeamRequest, + BudgetNewRequest, CommonProxyErrors, DeleteTeamRequest, LiteLLM_AccessGroupTable, @@ -57,6 +58,7 @@ from litellm.proxy._types import ( SpecialManagementEndpointEnums, SpecialModelNames, SpecialProxyStrings, + TeamAccessGroupModelGrant, TeamAddMemberResponse, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, @@ -155,6 +157,15 @@ router: Final = APIRouter() _DbRecordT = TypeVar("_DbRecordT") +class _TeamIdKeyCount(TypedDict): + team_id: int + + +class _TeamIdGroupRow(TypedDict): + team_id: str + _count: _TeamIdKeyCount + + class _PrismaTableActions(Protocol[_DbRecordT]): async def find_unique( self, @@ -219,59 +230,127 @@ class _PrismaTableActions(Protocol[_DbRecordT]): where: Mapping[str, object] | None = None, ) -> int: ... + async def group_by( + self, + by: Sequence[str], + where: Mapping[str, object] | None = None, + count: Mapping[str, bool] | None = None, + ) -> Sequence[_TeamIdGroupRow]: ... + + +class _HasTableActions(Protocol[_DbRecordT]): + @property + def table(self) -> "_PrismaTableActions[_DbRecordT]": ... + + +def _typed_table( + repo: "_HasTableActions[_DbRecordT]", record_type: type[_DbRecordT] +) -> "_PrismaTableActions[_DbRecordT]": + return repo.table + + +def _as_object(value: object) -> object: + return value + + +def _nullable(value: _DbRecordT | None) -> _DbRecordT | None: + return value + + +class _UserIdRow(Protocol): + @property + def user_id(self) -> str | None: ... + + +class _HasUserIdTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_UserIdRow]": ... + + +def _user_id_rows_db(repo: "_HasUserIdTable") -> "_PrismaTableActions[_UserIdRow]": + return repo.table + + +class _RawTeamRow(Protocol): + @property + def members_with_roles(self) -> Sequence[Mapping[str, object]] | None: ... + + +class _HasRawTeamTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_RawTeamRow]": ... + + +def _raw_team_db(repo: "_HasRawTeamTable") -> "_PrismaTableActions[_RawTeamRow]": + return repo.table + + +class _BudgetWriteCall(Protocol): + async def __call__( + self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth + ) -> LiteLLM_BudgetTableFull: ... + + +def _as_budget_write(fn: "_BudgetWriteCall") -> "_BudgetWriteCall": + return fn + + +class _TeamFindManyArgs(TypedDict, total=False): + take: int + skip: int + order: Mapping[str, str] + cursor: Mapping[str, object] + + +class _TeamUiViewFilters(TypedDict, total=False): + team_id: Mapping[str, str] + team_alias: Mapping[str, str] + + +class _TeamIdInFilter(TypedDict, total=False): + team_id: Mapping[str, Sequence[str]] + def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]": - team_table: Final[_PrismaTableActions[LiteLLM_TeamTable]] = TeamRepository(prisma_client).table - return team_table + return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable) def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]": - membership_table: Final[_PrismaTableActions[LiteLLM_TeamMembership]] = TeamMembershipRepository(prisma_client).table - return membership_table + return _typed_table(TeamMembershipRepository(prisma_client), LiteLLM_TeamMembership) def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]": - user_table: Final[_PrismaTableActions[LiteLLM_UserTable]] = UserRepository(prisma_client).table - return user_table + return _typed_table(UserRepository(prisma_client), LiteLLM_UserTable) def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]": - model_table: Final[_PrismaTableActions[LiteLLM_ModelTable]] = ModelTableRepository(prisma_client).table - return model_table + return _typed_table(ModelTableRepository(prisma_client), LiteLLM_ModelTable) def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]": - org_table: Final[_PrismaTableActions[LiteLLM_OrganizationTable]] = OrganizationRepository(prisma_client).table - return org_table + return _typed_table(OrganizationRepository(prisma_client), LiteLLM_OrganizationTable) def _org_membership_db( prisma_client: PrismaClient | None, ) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]": - org_membership_table: _PrismaTableActions[LiteLLM_OrganizationMembershipTable] = OrganizationMembershipRepository( - prisma_client - ).table - return org_membership_table + return _typed_table(OrganizationMembershipRepository(prisma_client), LiteLLM_OrganizationMembershipTable) def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]": - budget_table: Final[_PrismaTableActions[LiteLLM_BudgetTableFull]] = BudgetRepository(prisma_client).table - return budget_table + return _typed_table(BudgetRepository(prisma_client), LiteLLM_BudgetTableFull) def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]": - deleted_team_table: _PrismaTableActions[LiteLLM_DeletedTeamTable] = DeletedTeamRepository(prisma_client).table - return deleted_team_table + return _typed_table(DeletedTeamRepository(prisma_client), LiteLLM_DeletedTeamTable) def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]": - access_group_table: _PrismaTableActions[LiteLLM_AccessGroupTable] = AccessGroupRepository(prisma_client).table - return access_group_table + return _typed_table(AccessGroupRepository(prisma_client), LiteLLM_AccessGroupTable) def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]": - tokens_table: _PrismaTableActions[LiteLLM_VerificationToken] = VerificationTokenRepository(prisma_client).table - return tokens_table + return _typed_table(VerificationTokenRepository(prisma_client), LiteLLM_VerificationToken) def _sanitize_for_log(value: object) -> str: @@ -407,7 +486,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - team_member_budget_table: Final = await new_budget( + team_member_budget_table: Final = await _as_budget_write(new_budget)( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -455,7 +534,7 @@ class TeamMemberBudgetHandler: if team_member_budget_duration is not None: budget_request.budget_duration = team_member_budget_duration - budget_row: Final = await update_budget( + budget_row: Final = await _as_budget_write(update_budget)( budget_obj=budget_request, user_api_key_dict=user_api_key_dict, ) @@ -570,7 +649,7 @@ class TeamMemberBudgetHandler: ) if missing: - await TeamMembershipRepository(prisma_client).table.create_many( + await _team_membership_db(prisma_client).create_many( data=missing, skip_duplicates=True, # safety net against concurrent races ) @@ -1406,9 +1485,10 @@ async def new_team( complete_team_data_dict["metadata"] = encrypt_callback_vars(complete_team_data_dict["metadata"]) complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict) + team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict - team_row: Final[LiteLLM_TeamTable] = await TeamRepository(prisma_client).table.create( - data=complete_team_data_dict, + team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create( + data=team_creation_data, include={"litellm_model_table": True}, ) @@ -1855,7 +1935,7 @@ async def update_team( detail={"error": f"soft_budget must be a non-negative finite number. Received: {data.soft_budget}"}, ) - existing_team_row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": data.team_id}) + existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id}) if existing_team_row is None: raise HTTPException( @@ -1883,7 +1963,7 @@ async def update_team( ) if data.max_budget is not None: - existing_soft_budget: Final = getattr(existing_team_row, "soft_budget", None) + existing_soft_budget: Final[object] = _as_object(getattr(existing_team_row, "soft_budget", None)) soft_budget_to_check: Final = data.soft_budget if data.soft_budget is not None else existing_soft_budget if soft_budget_to_check is not None and isinstance(soft_budget_to_check, (int, float)): if data.max_budget <= soft_budget_to_check: @@ -1942,7 +2022,7 @@ async def update_team( data.organization_id = None # check org team limits - if updating team that belongs to an org - org_id_to_check: Final = ( + org_id_to_check: Final[object] = _as_object( data.organization_id if data.organization_id is not None else existing_team_row.organization_id ) if org_id_to_check is not None and isinstance(org_id_to_check, str) and prisma_client is not None: @@ -1975,7 +2055,7 @@ async def update_team( TeamMemberBudgetHandler.strip_system_managed_metadata_keys(updated_kv["metadata"]) if "metadata" in updated_kv: - stored_metadata: Final = ( + stored_metadata: Final[Mapping[str, JsonValue] | None] = ( { # mutable-ok: the validator payload's isinstance guard requires a plain dict key: value for key, value in existing_team_row.metadata.items() @@ -2078,16 +2158,19 @@ async def update_team( updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Final[LiteLLM_TeamTable | None] = await TeamRepository(prisma_client).table.update( - where={"team_id": data.team_id}, - data=updated_kv, - # `object_permission` is included so `_refresh_cached_team` - # doesn't write a cached team with the relation nulled out — - # see team_model_add for the full rationale. - include={ - "litellm_model_table": True, - "object_permission": True, - }, + team_update_data: Final[Mapping[str, object]] = updated_kv + team_row: Final[LiteLLM_TeamTable | None] = _nullable( + await _team_db(prisma_client).update( + where={"team_id": data.team_id}, + data=team_update_data, + # `object_permission` is included so `_refresh_cached_team` + # doesn't write a cached team with the relation nulled out — + # see team_model_add for the full rationale. + include={ + "litellm_model_table": True, + "object_permission": True, + }, + ) ) if team_row is None or team_row.team_id is None: @@ -2602,7 +2685,7 @@ async def _resolve_existing_member_user_ids( if not requested_user_ids: return frozenset() - found: Final = await UserRepository(prisma_client).table.find_many( + found: Final = await _user_id_rows_db(UserRepository(prisma_client)).find_many( where={ # mutable-ok: Prisma query filters are dict-shaped "user_id": { # mutable-ok: Prisma query filters are dict-shaped "in": sorted(requested_user_ids) @@ -3097,7 +3180,9 @@ async def team_member_delete( key_val["user_id"] = data.user_id elif data.user_email is not None: key_val["user_email"] = data.user_email - existing_user_rows: Final = await UserRepository(prisma_client).table.find_many(where=key_val) + existing_user_rows: Final[Sequence[LiteLLM_UserTable] | None] = await UserRepository(prisma_client).table.find_many( + where=key_val + ) if existing_user_rows is not None and (isinstance(existing_user_rows, list) and len(existing_user_rows) > 0): for existing_user in existing_user_rows: @@ -3105,7 +3190,7 @@ async def team_member_delete( if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - await UserRepository(prisma_client).table.update( + await _user_db(prisma_client).update( where={ "user_id": existing_user.user_id, }, @@ -3113,7 +3198,7 @@ async def team_member_delete( ) # Also clean up any existing team membership rows for this user and team - user_ids_to_delete: Final = set() + user_ids_to_delete: Final = set[str]() if data.user_id is not None: user_ids_to_delete.add(data.user_id) if existing_user_rows is not None and isinstance(existing_user_rows, list): @@ -3122,9 +3207,7 @@ async def team_member_delete( user_ids_to_delete.add(existing_user.user_id) for _uid in user_ids_to_delete: - await TeamMembershipRepository(prisma_client).table.delete_many( - where={"team_id": data.team_id, "user_id": _uid} - ) + await _team_membership_db(prisma_client).delete_many(where={"team_id": data.team_id, "user_id": _uid}) ## DELETE KEYS CREATED BY USER FOR THIS TEAM if user_ids_to_delete: @@ -3133,9 +3216,7 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await VerificationTokenRepository( - prisma_client - ).table.find_many( + keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3150,7 +3231,7 @@ async def team_member_delete( litellm_changed_by=None, ) - await VerificationTokenRepository(prisma_client).table.delete_many( + await _tokens_db(prisma_client).delete_many( where={ "user_id": {"in": list(user_ids_to_delete)}, "team_id": data.team_id, @@ -3310,7 +3391,7 @@ async def team_member_update( ### upsert new budget budget_patch: Final = _build_member_budget_patch(data) - async with prisma_client.db.tx() as tx: + async with prisma_client.tx() as tx: await _upsert_budget_and_membership( tx=tx, team_id=data.team_id, @@ -3653,7 +3734,7 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: list[LiteLLM_VerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many( + keys_to_delete: list[LiteLLM_VerificationToken] = await _tokens_db(prisma_client).find_many( where={"team_id": {"in": data.team_ids}} ) @@ -3829,7 +3910,7 @@ async def _add_team_member_budget_table( ) -> TeamInfoResponseObjectTeamTable: try: team_budget: Final = await _budget_db(prisma_client).find_unique(where={"budget_id": team_member_budget_id}) - team_info_response_object.team_member_budget_table = team_budget + return team_info_response_object.model_copy(update={"team_member_budget_table": team_budget}) except Exception: verbose_proxy_logger.info( "Team member budget table not found, passed team_member_budget_id=%s", team_member_budget_id @@ -3838,21 +3919,34 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _resolve_team_access_group_resources(_team_info: TeamInfoResponseObjectTeamTable) -> None: - """Populate access_group_models / mcp_server_ids / agent_ids on the team - info response by resolving inherited resources from its access groups.""" +async def _resolve_team_access_group_resources( + _team_info: TeamInfoResponseObjectTeamTable, +) -> TeamInfoResponseObjectTeamTable: + """Return a copy of the team info with access_group_models / mcp_server_ids / + agent_ids / details resolved from its access groups.""" if not _team_info.access_group_ids: - return + return _team_info ag_lookup: Final = await _batch_resolve_access_group_resources(_team_info.access_group_ids) - models, mcp_ids, agent_ids = set(), set(), set() - for ag_id in _team_info.access_group_ids: - if ag_id in ag_lookup: - models.update(ag_lookup[ag_id]["models"]) - mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) - agent_ids.update(ag_lookup[ag_id]["agent_ids"]) - _team_info.access_group_models = list(models) - _team_info.access_group_mcp_server_ids = list(mcp_ids) - _team_info.access_group_agent_ids = list(agent_ids) + resolved_groups: Final = tuple( + ag_lookup[ag_id] for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id in ag_lookup + ) + return _team_info.model_copy( + update={ + "access_group_models": list({m for group in resolved_groups for m in (group.access_model_names or [])}), + "access_group_mcp_server_ids": list( + {s for group in resolved_groups for s in (group.access_mcp_server_ids or [])} + ), + "access_group_agent_ids": list({a for group in resolved_groups for a in (group.access_agent_ids or [])}), + "access_group_details": tuple( + TeamAccessGroupModelGrant( + access_group_id=group.access_group_id, + access_group_name=group.access_group_name, + models=tuple(group.access_model_names or ()), + ) + for group in resolved_groups + ), + } + ) @router.get("/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -3958,11 +4052,11 @@ async def team_info( ) # Resolve resources inherited from access groups - await _resolve_team_access_group_resources(_team_info) + resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) response_object: Final = TeamInfoResponseObject( team_id=team_id, - team_info=_team_info, + team_info=resolved_team_info, keys=keys, team_memberships=returned_tm, ) @@ -4391,32 +4485,21 @@ async def _build_team_list_where_conditions( async def _batch_resolve_access_group_resources( all_access_group_ids: list[str], -) -> dict[str, dict[str, list[str]]]: +) -> dict[str, LiteLLM_AccessGroupTable]: """ - Batch-fetch access groups in a single DB query and return a per-group - resource map. - - Returns {ag_id: {"models": [...], "mcp_server_ids": [...], "agent_ids": [...]}}. - Missing/invalid groups are silently omitted. + Batch-fetch access groups in a single DB query and return them keyed by + access_group_id. Missing/invalid groups are silently omitted. """ from litellm.proxy.proxy_server import prisma_client as _prisma_client if not all_access_group_ids or _prisma_client is None: return {} - unique_ids: Final = list(set(all_access_group_ids)) + unique_ids: Final = tuple(frozenset(all_access_group_ids)) rows: Final = await _access_group_db(_prisma_client).find_many( where={"access_group_id": {"in": unique_ids}}, ) - - result: Final[dict[str, dict[str, list[str]]]] = {} - for row in rows: - result[row.access_group_id] = { - "models": list(row.access_model_names or []), - "mcp_server_ids": list(row.access_mcp_server_ids or []), - "agent_ids": list(row.access_agent_ids or []), - } - return result + return {row.access_group_id: row for row in rows} def _convert_teams_to_response_models( @@ -4466,7 +4549,7 @@ async def _get_keys_count_by_team( if not page_team_ids: return {} - grouped: Final = await VerificationTokenRepository(prisma_client).table.group_by( + grouped: Final = await _tokens_db(prisma_client).group_by( by=["team_id"], where={"team_id": {"in": page_team_ids}}, count={"team_id": True}, @@ -4710,15 +4793,18 @@ async def list_team_v2( all_ag_ids: Final = [ag_id for t in team_items_with_ag for ag_id in (t.access_group_ids or [])] ag_lookup: Final = await _batch_resolve_access_group_resources(all_ag_ids) for team_item in team_items_with_ag: - models, mcp_ids, agent_ids = set(), set(), set() - for ag_id in team_item.access_group_ids or []: - if ag_id in ag_lookup: - models.update(ag_lookup[ag_id]["models"]) - mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) - agent_ids.update(ag_lookup[ag_id]["agent_ids"]) - team_item.access_group_models = list(models) - team_item.access_group_mcp_server_ids = list(mcp_ids) - team_item.access_group_agent_ids = list(agent_ids) + team_groups = tuple( + ag_lookup[ag_id] for ag_id in (team_item.access_group_ids or []) if ag_id in ag_lookup + ) + team_item.access_group_models = list( + {m for group in team_groups for m in (group.access_model_names or [])} + ) + team_item.access_group_mcp_server_ids = list( + {s for group in team_groups for s in (group.access_mcp_server_ids or [])} + ) + team_item.access_group_agent_ids = list( + {a for group in team_groups for a in (group.access_agent_ids or [])} + ) return { "teams": team_list, @@ -4780,7 +4866,7 @@ async def _authorize_and_filter_teams( if allowed_org_ids is not None: # Org admin: query DB for teams in their orgs - org_teams: Final = await TeamRepository(prisma_client).table.find_many( + org_teams: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( where={"organization_id": {"in": allowed_org_ids}}, include={"litellm_model_table": True}, ) @@ -4794,7 +4880,9 @@ async def _authorize_and_filter_teams( ] elif user_id: # Regular user: fetch all and filter by membership (Prisma can't filter JSON arrays) - response: Final = await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True}) + response: Final = await _raw_team_db(TeamRepository(prisma_client)).find_many( + include={"litellm_model_table": True} + ) return [ team for team in response @@ -4802,7 +4890,7 @@ async def _authorize_and_filter_teams( ] else: # Proxy admin: all teams - return list(await TeamRepository(prisma_client).table.find_many(include={"litellm_model_table": True})) + return list(await _raw_team_db(TeamRepository(prisma_client)).find_many(include={"litellm_model_table": True})) @router.get("/team/list", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -4854,7 +4942,7 @@ async def list_team( _team_memberships.append(tm) # add all keys that belong to the team - keys = await VerificationTokenRepository(prisma_client).table.find_many(where={"team_id": team.team_id}) + keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id}) try: returned_responses.append( @@ -4905,7 +4993,7 @@ async def get_paginated_teams( total_count: Final = await _team_db(prisma_client).count() # Get paginated teams - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( skip=skip, take=page_size, order={"team_alias": "asc"}, # Sort by team_alias @@ -4955,7 +5043,7 @@ async def ui_view_teams( skip: Final = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions: Final = {} + where_conditions: Final[_TeamUiViewFilters] = {} if team_id: where_conditions["team_id"] = { @@ -4970,7 +5058,7 @@ async def ui_view_teams( } # Query users with pagination and filters - teams: Final = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _team_db(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -5160,13 +5248,13 @@ async def team_model_delete( ) # Get current models list - current_models: Final = team_obj.models or [] + current_models: Final[Sequence[str]] = team_obj.models or [] # Remove specified models updated_models: Final = [m for m in current_models if m not in data.models] # Update team. See team_model_add for the rationale on `include`. - updated_team: Final = await TeamRepository(prisma_client).table.update( + updated_team: Final = await _team_db(prisma_client).update( where={"team_id": data.team_id}, data={"models": updated_models}, include={"object_permission": True}, @@ -5419,7 +5507,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi BATCH_SIZE: Final = 500 while True: - find_args: dict = { + find_args: _TeamFindManyArgs = { "take": BATCH_SIZE, "order": {"team_id": "asc"}, } @@ -5427,7 +5515,7 @@ async def _append_permissions_to_all_teams(prisma_client: PrismaClient, permissi find_args["cursor"] = {"team_id": cursor} find_args["skip"] = 1 - teams = await TeamRepository(prisma_client).table.find_many(**find_args) + teams = await _team_db(prisma_client).find_many(**find_args) if not teams: break @@ -5522,11 +5610,11 @@ async def get_team_daily_activity( ) ## Fetch team aliases and check team admin status - where_condition: Final = {} + where_condition: Final[_TeamIdInFilter] = {} if team_ids_list: where_condition["team_id"] = {"in": list(team_ids_list)} - team_aliases: Final = await TeamRepository(prisma_client).table.find_many(where=where_condition) - team_alias_metadata: Final = {t.team_id: {"team_alias": t.team_alias} for t in team_aliases} + team_aliases: Final = await _team_db(prisma_client).find_many(where=where_condition) + team_alias_metadata: Final = {t.team_id: {"team_alias": _as_object(t.team_alias)} for t in team_aliases} # Check if user is team admin or has /team/daily/activity permission # If not, filter by user's API keys. diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 44abc56713f..a2c50590dd5 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -16,9 +16,22 @@ import json import os import re import secrets +from collections.abc import Mapping, Sequence from copy import deepcopy from html import escape -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NoReturn, + Optional, + Protocol, + TypeVar, + Union, + cast, + overload, +) from urllib.parse import parse_qs, urlencode, urlparse if TYPE_CHECKING: @@ -155,6 +168,102 @@ _CLI_SSO_SECRET_KEY_FRAGMENTS: Final = frozenset( } ) +_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True) + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique( + self, + where: Mapping[str, object], + ) -> _DbRecordT | None: ... + + async def find_first( + self, + where: Mapping[str, object] | None = None, + ) -> _DbRecordT | None: ... + + async def find_many( + self, + where: Mapping[str, object] | None = None, + ) -> Sequence[_DbRecordT]: ... + + async def update( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _DbRecordT: ... + + async def update_many( + self, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> int: ... + + +class _UserMetadataRow(Protocol): + @property + def metadata(self) -> Mapping[str, object] | None: ... + + +class _HasUserMetadataTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_UserMetadataRow]": ... + + +def _user_meta_db(repo: "_HasUserMetadataTable") -> "_PrismaTableActions[_UserMetadataRow]": + return repo.table + + +class _SsoConfigRow(Protocol): + @property + def sso_settings(self) -> Mapping[str, object] | None: ... + + +class _HasSsoConfigTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_SsoConfigRow]": ... + + +def _sso_config_db(repo: "_HasSsoConfigTable") -> "_PrismaTableActions[_SsoConfigRow]": + return repo.table + + +class _TeamDetailRow(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _HasTeamDetailTable(Protocol): + @property + def table(self) -> "_PrismaTableActions[_TeamDetailRow]": ... + + +def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDetailRow]": + return repo.table + + +class _CustomSsoCall(Protocol): + async def __call__(self, sso_response: object) -> SSOUserDefinedValues | None: ... + + +class _ServicePrincipalAssignment(Protocol): + def get(self, key: str) -> str: ... + + +class _ServicePrincipalPage(Protocol): + @overload + def get( + self, + key: Literal["value"], + default: Sequence["_ServicePrincipalAssignment"], + ) -> Sequence["_ServicePrincipalAssignment"]: ... + + @overload + def get(self, key: Literal["@odata.nextLink"]) -> str | None: ... + + +def _as_object(value: object) -> object: + return value + def _hash_cli_sso_secret(secret: str) -> str: return hashlib.sha256(secret.encode("utf-8")).hexdigest() @@ -256,7 +365,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict: flow = cache.get_cache(key=cache_key) if isinstance(flow, str): try: - flow = json.loads(flow) + flow = _as_object(json.loads(flow)) except ValueError: flow = None if not isinstance(flow, dict) or "poll_secret_hash" not in flow: @@ -421,7 +530,7 @@ def _flatten_cli_sso_metadata_for_poll( def build_cli_sso_attribution_metadata( result: CustomOpenID | OpenID | dict, -) -> dict[str, Any]: +) -> dict[str, object]: """ Build allowlisted, non-secret scalar attribution metadata from an SSO result. @@ -432,7 +541,7 @@ def build_cli_sso_attribution_metadata( if not claim_map: return {} - metadata: Final[dict[str, Any]] = {} + metadata: Final[dict[str, object]] = {} for source_claim, dest_key in claim_map: if not _is_safe_cli_sso_metadata_dest_key(dest_key): verbose_proxy_logger.debug("Skipping unsafe CLI SSO metadata destination key: %s", dest_key) @@ -474,14 +583,14 @@ def _merge_cli_sso_attribution_metadata( async def _persist_cli_sso_user_metadata( prisma_client: PrismaClient, user_id: str, - attribution_metadata: dict[str, Any], + attribution_metadata: dict[str, object], ) -> None: if not attribution_metadata: return try: - user_row: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) - existing_metadata: dict[str, Any] = {} + user_row: Final = await _user_meta_db(UserRepository(prisma_client)).find_unique(where={"user_id": user_id}) + existing_metadata: dict[str, object] = {} if user_row is not None: row_metadata: Final = user_row.metadata if isinstance(row_metadata, dict): @@ -491,7 +600,7 @@ async def _persist_cli_sso_user_metadata( existing_metadata=existing_metadata, attribution_metadata=attribution_metadata, ) - await UserRepository(prisma_client).table.update_many( + await _user_meta_db(UserRepository(prisma_client)).update_many( where={"user_id": user_id}, data={"metadata": merged_metadata}, ) @@ -1104,7 +1213,7 @@ def generic_response_convertor( ) # Build extra_fields dict from GENERIC_USER_EXTRA_ATTRIBUTES if specified - extra_fields: dict[str, Any] | None = None + extra_fields: dict[str, object] | None = None if generic_user_extra_attributes: extra_fields = {} for attr_name in generic_user_extra_attributes.split(","): @@ -1193,7 +1302,9 @@ async def _setup_team_mappings() -> Optional["TeamMappings"]: prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict: Final = dict(sso_db_record.sso_settings) @@ -1225,7 +1336,9 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: prisma_client: Final = get_prisma_client_or_throw("Prisma client is None, connect a database to your proxy") - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_config_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) if sso_db_record and sso_db_record.sso_settings: sso_settings_dict: Final = dict(sso_db_record.sso_settings) @@ -1273,7 +1386,7 @@ async def _setup_role_mappings() -> Optional["RoleMappings"]: return role_mappings -def _parse_generic_sso_headers() -> dict: +def _parse_generic_sso_headers() -> dict[str, str]: """Parse comma-separated GENERIC_SSO_HEADERS env var into a dict.""" raw: Final = os.getenv("GENERIC_SSO_HEADERS", None) if raw is None: @@ -1677,7 +1790,7 @@ def _build_sso_user_update_data( result: Union["CustomOpenID", OpenID, dict] | None, user_email: str | None, user_id: str | None, -) -> dict: +) -> dict[str, object]: """ Build the update data dictionary for SSO user upsert. @@ -1689,7 +1802,7 @@ def _build_sso_user_update_data( Returns: dict: Update data containing user_email and optionally user_role if valid """ - update_data: Final[dict] = {"user_email": normalize_email(user_email)} + update_data: Final[dict[str, object]] = {"user_email": normalize_email(user_email)} # Get SSO role from result and include if valid sso_role: Final = getattr(result, "user_role", None) @@ -1740,7 +1853,7 @@ async def _sync_user_role_from_jwt_role_map( # Update existing DB record if role differs if user_info is not None and user_info.user_role != mapped_role.value: - await UserRepository(prisma_client).table.update( + await _user_meta_db(UserRepository(prisma_client)).update( where={"user_id": user_info.user_id}, data={"user_role": mapped_role.value}, ) @@ -1796,7 +1909,7 @@ async def check_and_update_if_proxy_admin_id(user_role: str, user_id: str, prism return user_role if prisma_client: - await UserRepository(prisma_client).table.update( + await _user_meta_db(UserRepository(prisma_client)).update( where={"user_id": user_id}, data={"user_role": LitellmUserRoles.PROXY_ADMIN.value}, ) @@ -2016,10 +2129,11 @@ async def _build_cli_sso_user_defined_values( ) -> SSOUserDefinedValues | None: from litellm.proxy.proxy_server import user_custom_sso + custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso user_id: Final = parsed_openid_result.get("user_id") - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - return await user_custom_sso(result) + if custom_sso_handler is not None: + if inspect.iscoroutinefunction(custom_sso_handler): + return await custom_sso_handler(result) raise ValueError("user_custom_sso must be a coroutine function") if user_id is None: return None @@ -2035,12 +2149,14 @@ async def _build_cli_sso_user_defined_values( async def _fetch_cli_sso_team_details( prisma_client: PrismaClient, - teams: list[str], -) -> list[dict[str, Any]]: - team_details: Final[list[dict[str, Any]]] = [] + teams: Sequence[str], +) -> list[dict[str, object]]: + team_details: Final[list[dict[str, object]]] = [] try: if teams: - prisma_teams: Final = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": teams}}) + prisma_teams: Final = await _team_detail_db(TeamRepository(prisma_client)).find_many( + where={"team_id": {"in": teams}} + ) for team_row in prisma_teams: team_dict = team_row.model_dump() team_details.append( @@ -2257,12 +2373,12 @@ async def cli_poll_key( verbose_proxy_logger.info("Returning teams list for user %s to select from: %s", user_id, user_teams) # Best-effort construction of team_details if it wasn't # already cached for some reason. - team_details_response: list[dict[str, Any]] | None = None + team_details_response: list[dict[str, object]] | None = None if isinstance(user_team_details, list) and user_team_details: team_details_response = user_team_details elif user_teams: team_details_response = [{"team_id": t, "team_alias": None} for t in user_teams] - poll_response: dict[str, Any] = { + poll_response: dict[str, object] = { "status": "ready", "user_id": user_id, "teams": user_teams, @@ -2997,7 +3113,9 @@ class SSOAuthenticationHandler: user_id=user_id, ) - await UserRepository(prisma_client).table.update_many(where={"user_id": user_id}, data=update_data) + await _user_meta_db(UserRepository(prisma_client)).update_many( + where={"user_id": user_id}, data=update_data + ) else: verbose_proxy_logger.info("user not in DB, inserting user into LiteLLM DB") # user not in DB, insert User into LiteLLM DB @@ -3089,7 +3207,9 @@ class SSOAuthenticationHandler: code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) try: - team_obj: Final = await TeamRepository(prisma_client).table.find_first(where={"team_id": litellm_team_id}) + team_obj: Final = await _team_detail_db(TeamRepository(prisma_client)).find_first( + where={"team_id": litellm_team_id} + ) verbose_proxy_logger.debug("Team object: %s", team_obj) # only create a new team if it doesn't exist @@ -3278,9 +3398,10 @@ class SSOAuthenticationHandler: # But if it is, we want their models preferences user_defined_values: SSOUserDefinedValues | None = None - if user_custom_sso is not None: - if inspect.iscoroutinefunction(user_custom_sso): - user_defined_values = await user_custom_sso(result) + custom_sso_handler: Final[_CustomSsoCall | None] = user_custom_sso + if custom_sso_handler is not None: + if inspect.iscoroutinefunction(custom_sso_handler): + user_defined_values = await custom_sso_handler(result) else: raise ValueError("user_custom_sso must be a coroutine function") elif user_id is not None: @@ -3448,7 +3569,7 @@ class SSOAuthenticationHandler: dict: Token exchange parameters """ # Prepare token exchange parameters (may add code_verifier: str later) - token_params: Final[dict[str, Any]] = {"include_client_id": generic_include_client_id} + token_params: Final[dict[str, object]] = {"include_client_id": generic_include_client_id} # Retrieve PKCE code_verifier if PKCE was used in authorization. # Gate on GENERIC_CLIENT_USE_PKCE to avoid an unnecessary Redis round-trip @@ -3663,7 +3784,7 @@ class SSOAuthenticationHandler: access_token string. Raises ProxyException on any validation failure. """ try: - token_response_raw: Final = response.json() + token_response_raw: Final[object] = _as_object(response.json()) except Exception as json_err: verbose_proxy_logger.error( "Failed to parse token response as JSON: %s. Body: %s", @@ -4253,7 +4374,7 @@ class MicrosoftSSOHandler: while next_link is not None and page_count < MicrosoftSSOHandler.MAX_GRAPH_API_PAGES: response = await async_client.get(next_link, headers=headers) - response_json = response.json() + response_json: _ServicePrincipalPage = response.json() verbose_proxy_logger.debug("Response from service principal app role assigned to: %s", response_json) for _object in response_json.get("value", []): diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index fce57e514c1..56e986c89cf 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1,9 +1,10 @@ import base64 import mimetypes import re +from collections.abc import Mapping from dataclasses import dataclass, field from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, runtime_checkable from litellm.repositories.table_repositories import ( ManagedFileRepository, @@ -16,10 +17,26 @@ if TYPE_CHECKING: from prisma.models import LiteLLM_ManagedObjectTable from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.types.utils import LiteLLMBatch +@runtime_checkable +class ManagedResourceAccessChecker(Protocol): + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> bool: ... + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: "UserAPIKeyAuth", + ) -> bool: ... + + def _is_base64_encoded_unified_file_id(b64_uid: str) -> str | Literal[False]: # Ensure b64_uid is a string and not a mock object if not isinstance(b64_uid, str): @@ -879,6 +896,65 @@ def validate_managed_files_requirement( ) +async def validate_managed_id_requirement( + resource_id: str | None, + resource_kind: Literal["file", "batch", "fine-tuning job"], + user_api_key_dict: "UserAPIKeyAuth", + managed_files_obj: object | None, +) -> None: + """ + Enforce proxy-level managed resources on every route that accepts a provider-issued id + when ``litellm.require_managed_files`` is enabled, and authenticate managed ids against + the caller's stored ownership record. + + Ownership is only recorded for LiteLLM managed ids, so a raw provider id is forwarded to the + provider under shared credentials without any tenant check; knowing another tenant's provider + id would be enough to read, reuse, or destroy the object behind it. + + Raises: + HTTPException: 400 for a raw id, 403 for an inaccessible managed id, or 500 when + ownership validation is unavailable. + """ + from fastapi import HTTPException + + import litellm + + if litellm.require_managed_files is not True: + return + + if not resource_id: + return + + if not _is_base64_encoded_unified_file_id(resource_id): + raise HTTPException( + status_code=400, + detail=( + f"Raw provider {resource_kind} ids cannot be used when require_managed_files is enabled in " + f"litellm_settings. Use the LiteLLM managed {resource_kind} id returned when the " + f"{resource_kind} was created." + ), + ) + + if not isinstance(managed_files_obj, ManagedResourceAccessChecker): + raise HTTPException( + status_code=500, + detail="Managed resource ownership validation is unavailable.", + ) + + can_access: Final = ( + await managed_files_obj.can_user_call_unified_file_id(resource_id, user_api_key_dict) + if resource_kind == "file" + else await managed_files_obj.can_user_call_unified_object_id(resource_id, user_api_key_dict) + ) + if can_access: + return + + raise HTTPException( + status_code=403, + detail=f"The caller does not have access to this managed {resource_kind} id.", + ) + + def _extract_model_param(request: "Request", request_body: dict) -> str | None: """ Extract model parameter from request. @@ -1002,6 +1078,34 @@ async def resolve_output_file_ids_to_unified(response, prisma_client) -> None: pass +async def map_raw_file_ids_to_unified( + raw_file_ids: frozenset[str], prisma_client: "PrismaClient | None" +) -> Mapping[str, str]: + if not raw_file_ids or not prisma_client: + return MappingProxyType({}) + managed_files: Final = await ManagedFileRepository(prisma_client).table.find_many( + where={"flat_model_file_ids": {"hasSome": sorted(raw_file_ids)}} # mutable-ok: prisma where is a plain dict + ) + return MappingProxyType( + { + raw_id: managed_file.unified_file_id + for managed_file in managed_files + for raw_id in managed_file.flat_model_file_ids + if raw_id in raw_file_ids + } + ) + + +def apply_unified_file_ids(response: "LiteLLMBatch", unified_id_by_raw_id: Mapping[str, str]) -> None: + for file_attr, raw_id in ( + ("input_file_id", getattr(response, "input_file_id", None)), + ("output_file_id", getattr(response, "output_file_id", None)), + ("error_file_id", getattr(response, "error_file_id", None)), + ): + if isinstance(raw_id, str) and raw_id in unified_id_by_raw_id: + setattr(response, file_attr, unified_id_by_raw_id[raw_id]) + + async def ensure_batch_response_managed_file_ids( response, managed_files_obj, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index bf7aa96121a..0acaac3bf5d 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -50,6 +50,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( handle_model_based_routing, prepare_data_with_credentials, validate_managed_files_requirement, + validate_managed_id_requirement, ) from litellm.proxy.utils import ProxyLogging, is_known_model from litellm.repositories.table_repositories import ManagedFileRepository @@ -612,6 +613,13 @@ async def get_file_content( data: dict = {"file_id": file_id} try: + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + # Include original request and headers in the data base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) ( @@ -908,6 +916,13 @@ async def get_file( data: dict = {"file_id": file_id} try: + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + custom_llm_provider: Final = ( provider or get_custom_llm_provider_from_request_headers(request=request) @@ -1098,6 +1113,13 @@ async def delete_file( data: dict = {"file_id": file_id} try: + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + ) + custom_llm_provider: Final = ( provider or get_custom_llm_provider_from_request_headers(request=request) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 64e52d252ca..8a526fcd6cb 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -565,6 +565,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): # real parent span. _metadata["user_api_key"] = user_api_key_dict.api_key _metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span + _metadata.update( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict) + ) kwargs: Final = { "litellm_params": { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 07eaed9fe45..90daaaeae6b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17,7 +17,7 @@ import traceback import warnings from collections.abc import AsyncGenerator, Callable, Mapping from datetime import datetime, timedelta, timezone -from types import UnionType +from types import MappingProxyType, UnionType from typing import ( TYPE_CHECKING, Any, @@ -297,6 +297,9 @@ from litellm.proxy.common_request_processing import ( _should_return_raw_model_name, create_response, ) +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AuthCacheInvalidationSubscriber, +) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers @@ -868,7 +871,7 @@ async def proxy_shutdown_event(): async def _initialize_shared_aiohttp_session(): """Initialize shared aiohttp session for connection reuse with connection limits.""" try: - from aiohttp import ClientSession, TCPConnector + from aiohttp import ClientSession, DummyCookieJar, TCPConnector from litellm.llms.custom_httpx.http_handler import ( _build_aiohttp_keepalive_socket_factory, @@ -889,7 +892,7 @@ async def _initialize_shared_aiohttp_session(): connector_kwargs["socket_factory"] = socket_factory connector: Final = TCPConnector(**connector_kwargs) - session: Final = ClientSession(connector=connector) + session: Final = ClientSession(connector=connector, cookie_jar=DummyCookieJar()) verbose_proxy_logger.info( "SESSION REUSE: Created shared aiohttp session for connection pooling (ID: %s, limit=%s, limit_per_host=%s)", @@ -1013,6 +1016,37 @@ async def proxy_startup_event(app: FastAPI): asyncio.create_task(_run_pw_migration()) + async def _run_agent_grant_id_migration() -> None: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + object_permission_table, + ) + + for attempt in range(3): + try: + result = await global_agent_registry.migrate_legacy_grant_ids( + table=object_permission_table(prisma_client) + ) + if result.rewritten: + verbose_proxy_logger.info( + "Rewrote %s object_permission rows from legacy config agent ids", result.rewritten + ) + if result.missed == 0: + return + verbose_proxy_logger.warning( + "Legacy agent grant id migration attempt %s/3 left %s rows unmigrated", + attempt + 1, + result.missed, + ) + except Exception as e: # noqa: BLE001 # startup task must survive any DB error and retry + verbose_proxy_logger.warning( + "Legacy agent grant id migration attempt %s/3 failed: %s", attempt + 1, e + ) + if attempt < 2: + await asyncio.sleep(5) + + asyncio.create_task(_run_agent_grant_id_migration()) + ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before ## the coordination Redis is published to its consumers below. @@ -1111,6 +1145,10 @@ async def proxy_startup_event(app: FastAPI): prisma_client=prisma_client, ) ) + ProxyStartupEvent._warn_budget_without_db( + max_budget=litellm.max_budget, + prisma_client=prisma_client, + ) ### START BATCH WRITING DB + CHECKING NEW MODELS### if prisma_client is not None: @@ -1189,6 +1227,8 @@ async def proxy_startup_event(app: FastAPI): await proxy_config.stop_config_sync_subscriber() + await proxy_config.stop_auth_cache_invalidation_subscriber() + await proxy_shutdown_event() @@ -3900,6 +3940,7 @@ class ProxyConfig: self._last_hashicorp_vault_config: dict[str, Any] | None = None self.worker_registry: list[WorkerRegistryEntry] = [] self.config_sync_subscriber: ConfigSyncSubscriber | None = None + self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None from litellm.litellm_core_utils.get_model_cost_map import ( get_model_cost_map_loaded_at, ) @@ -6086,6 +6127,15 @@ class ProxyConfig: else: general_settings["disable_auto_add_proxy_admin_to_teams"] = value if value is None else bool(value) + if "apply_user_budget_to_team_keys" in _general_settings and ( + "apply_user_budget_to_team_keys" not in self._yaml_general_settings_keys + ): + db_value: Final = _general_settings["apply_user_budget_to_team_keys"] + if isinstance(db_value, str): + general_settings["apply_user_budget_to_team_keys"] = db_value.lower() == "true" + else: + general_settings["apply_user_budget_to_team_keys"] = db_value if db_value is None else bool(db_value) + ## STORE MODEL IN DB ## if "store_model_in_db" in _general_settings: value = _general_settings["store_model_in_db"] @@ -6387,6 +6437,30 @@ class ProxyConfig: except Exception as e: verbose_proxy_logger.error("Error stopping config sync subscriber: %s", e) + def start_auth_cache_invalidation_subscriber( + self, + redis_cache: RedisCache | None, + user_api_key_cache: UserApiKeyCache, + ) -> None: + if redis_cache is None or self.auth_cache_invalidation_subscriber is not None: + return + subscriber: Final = AuthCacheInvalidationSubscriber( + redis_cache=redis_cache, + user_api_key_cache=user_api_key_cache, + ) + self.auth_cache_invalidation_subscriber = subscriber + subscriber.start() + + async def stop_auth_cache_invalidation_subscriber(self) -> None: + subscriber: Final = self.auth_cache_invalidation_subscriber + if subscriber is None: + return + self.auth_cache_invalidation_subscriber = None + try: + await subscriber.stop() + except Exception as e: # noqa: BLE001 # best-effort: a failing stop must not break proxy shutdown + verbose_proxy_logger.error("Error stopping auth cache invalidation subscriber: %s", e) + async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): """ Use this to read non-llm objects from the db and initialize them @@ -7825,6 +7899,19 @@ def giveup(e): class ProxyStartupEvent: + @staticmethod + def _warn_budget_without_db(max_budget: float | None, prisma_client: PrismaClient | None) -> None: + if prisma_client is not None or not max_budget or max_budget <= 0: + return + + verbose_proxy_logger.warning( + "A proxy-wide budget (litellm.max_budget=%s) is configured but no database is connected, " + "so the budget will NOT be enforced and requests will never be blocked. Set DATABASE_URL or " + "general_settings.database_url and restart. Redis and fail_closed_budget_enforcement do not " + "cover the proxy-wide budget because there is no global spend counter; Redis alone is not a substitute.", + max_budget, + ) + @classmethod def _initialize_startup_logging( cls, @@ -8313,6 +8400,11 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + proxy_config.start_auth_cache_invalidation_subscriber( + redis_cache=redis_usage_cache, + user_api_key_cache=user_api_key_cache, + ) + if store_model_in_db is True: ### GET STORED CREDENTIALS ### scheduler.add_job( @@ -14930,6 +15022,29 @@ Keep it more precise, to prevent overwrite other values unintentially _PLUGIN_KEY_REDACTED: Final = "***" +_GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingProxyType( + { + "max_parallel_requests": "Integer", + "global_max_parallel_requests": "Integer", + "max_request_size_mb": "Integer", + "max_response_size_mb": "Integer", + "proxy_config_reload_interval_seconds": "Integer", + "pass_through_endpoints": "PydanticModel", + "store_model_in_db": "Boolean", + "store_prompts_in_spend_logs": "Boolean", + "maximum_spend_logs_retention_period": "String", + "mcp_internal_ip_ranges": "List", + "mcp_trusted_proxy_ranges": "List", + "mcp_xff_num_trusted_hops": "Integer", + "always_include_stream_usage": "Boolean", + "forward_client_headers_to_llm_api": "Boolean", + "mcp_required_fields": "List", + "cancel_on_disconnect": "Boolean", + "disable_auto_add_proxy_admin_to_teams": "Boolean", + "apply_user_budget_to_team_keys": "Boolean", + } +) + def _preserve_redacted_plugin_keys(incoming: object, existing: object) -> object: """Restore real plugin_key values the client never sees. @@ -15428,25 +15543,7 @@ async def get_config_list( else: db_general_settings_dict = {} - allowed_args: Final = { - "max_parallel_requests": {"type": "Integer"}, - "global_max_parallel_requests": {"type": "Integer"}, - "max_request_size_mb": {"type": "Integer"}, - "max_response_size_mb": {"type": "Integer"}, - "proxy_config_reload_interval_seconds": {"type": "Integer"}, - "pass_through_endpoints": {"type": "PydanticModel"}, - "store_model_in_db": {"type": "Boolean"}, - "store_prompts_in_spend_logs": {"type": "Boolean"}, - "maximum_spend_logs_retention_period": {"type": "String"}, - "mcp_internal_ip_ranges": {"type": "List"}, - "mcp_trusted_proxy_ranges": {"type": "List"}, - "mcp_xff_num_trusted_hops": {"type": "Integer"}, - "always_include_stream_usage": {"type": "Boolean"}, - "forward_client_headers_to_llm_api": {"type": "Boolean"}, - "mcp_required_fields": {"type": "List"}, - "cancel_on_disconnect": {"type": "Boolean"}, - "disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"}, - } + allowed_args: Final = _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES return_val: Final = [] @@ -15454,7 +15551,7 @@ async def get_config_list( if field_name in allowed_args: ## HANDLE TYPED DICT - typed_dict_type = allowed_args[field_name]["type"] + typed_dict_type = allowed_args[field_name] if typed_dict_type == "PydanticModel": if field_name == "pass_through_endpoints": @@ -15496,7 +15593,7 @@ async def get_config_list( _response_obj = ConfigList( field_name=field_name, - field_type=allowed_args[field_name]["type"], + field_type=allowed_args[field_name], field_description=field_info.description or "", field_value=_redact_general_setting_value( field_name, @@ -15524,7 +15621,7 @@ async def get_config_list( _response_obj = ConfigList( field_name=field_name, - field_type=allowed_args[field_name]["type"], + field_type=allowed_args[field_name], field_description=field_info.description or "", field_value=_redact_general_setting_value(field_name, _field_value, is_full_admin), stored_in_db=_stored_in_db, diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index 78c3e9fd31b..79791607b2e 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -220,7 +220,7 @@ async def get_agents(request: Request): "url": get_custom_url(str(request.base_url), route=f"a2a/{agent.agent_id}"), } for agent in agents - if agent.agent_id in litellm.public_agent_groups + if not global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups) ] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 907c4ac7344..58a85171cc7 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -156,6 +156,7 @@ async def reserve_budget_for_request( proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: Any | None = None, + apply_user_budget_to_team_keys: bool = False, fail_closed_budget_enforcement: bool = False, ) -> dict | None: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): @@ -175,6 +176,7 @@ async def reserve_budget_for_request( proxy_logging_obj=proxy_logging_obj, end_user_id=end_user_id, end_user_object=end_user_object, + apply_user_budget_to_team_keys=apply_user_budget_to_team_keys, ) if not counters: return None @@ -332,6 +334,7 @@ async def _get_budget_counters( proxy_logging_obj: ProxyLogging, end_user_id: str | None = None, end_user_object: Any | None = None, + apply_user_budget_to_team_keys: bool = False, ) -> list[_BudgetCounter]: counters: Final[list[_BudgetCounter]] = [] @@ -380,8 +383,9 @@ async def _get_budget_counters( ) ) + is_team_key: Final = team_object is not None and team_object.team_id is not None if ( - (team_object is None or team_object.team_id is None) + (not is_team_key or apply_user_budget_to_team_keys) and user_object is not None and user_object.user_id is not None and user_object.max_budget is not None diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 382df608a0c..08bb8698cac 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -4,11 +4,11 @@ import json import os from collections import Counter from collections.abc import Mapping -from typing import Any, Final +from typing import Any, Final, Protocol, TypeVar from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile -from pydantic import ConfigDict, ValidationError, create_model +from pydantic import ConfigDict, JsonValue, ValidationError, create_model from pydantic.fields import FieldInfo import litellm @@ -36,6 +36,73 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( router: Final = APIRouter() +_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True) + + +class _PrismaTableActions(Protocol[_DbRecordT]): + async def find_unique(self, where: Mapping[str, object]) -> _DbRecordT | None: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ... + + async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _DbRecordT: ... + + +class _SsoSettingsMappingRow(Protocol): + @property + def sso_settings(self) -> Mapping[str, object] | None: ... + + +class _HasSsoSettingsMappingTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_SsoSettingsMappingRow]: ... + + +def _sso_settings_mapping_db(repo: _HasSsoSettingsMappingTable) -> _PrismaTableActions[_SsoSettingsMappingRow]: + return repo.table + + +class _StoredSsoSettingsRow(Protocol): + @property + def sso_settings(self) -> object: ... + + +class _HasStoredSsoSettingsTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_StoredSsoSettingsRow]: ... + + +def _stored_sso_settings_db(repo: _HasStoredSsoSettingsTable) -> _PrismaTableActions[_StoredSsoSettingsRow]: + return repo.table + + +class _UiSettingsRow(Protocol): + @property + def ui_settings(self) -> str | Mapping[str, JsonValue] | None: ... + + +class _HasUiSettingsTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_UiSettingsRow]: ... + + +def _ui_settings_db(repo: _HasUiSettingsTable) -> _PrismaTableActions[_UiSettingsRow]: + return repo.table + + +class _ConfigParamRow(Protocol): + @property + def param_value(self) -> str | Mapping[str, object] | None: ... + + +class _HasConfigParamTable(Protocol): + @property + def table(self) -> _PrismaTableActions[_ConfigParamRow]: ... + + +def _config_param_db(repo: _HasConfigParamTable) -> _PrismaTableActions[_ConfigParamRow]: + return repo.table + + # Maps each UIThemeConfig field to the env var the UI branding path reads it # from. /update/ui_theme_settings writes both the stored ui_theme_config and # these env vars, so /get/ui_theme_settings resolves the same env vars to @@ -54,7 +121,7 @@ def _is_public_http_url(value: str | None) -> bool: return parsed.scheme in ("http", "https") and bool(parsed.netloc) -def _resolve_ui_theme_field(stored_values: Mapping[str, Any], field_name: str) -> str | None: +def _resolve_ui_theme_field(stored_values: Mapping[str, object], field_name: str) -> str | None: """Resolve one UI theme field to the value the branding path actually uses. The stored ui_theme_config wins; a field absent or blank there falls back to @@ -263,7 +330,7 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [ # include generics like ``Optional[int]`` / ``List[str]`` that are not # instances of ``type`` — so tightening this to ``type`` would reject # valid inputs. -_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[Any, FieldInfo]]] = {} +_EXTRA_UI_SETTINGS_FIELDS: Final[dict[str, tuple[object, FieldInfo]]] = {} # Settings OSS knows about as enterprise-gated. If a caller sends one of # these keys and no extension package has registered it, the PATCH @@ -275,7 +342,7 @@ _ENTERPRISE_ONLY_UI_SETTINGS: Final[set[str]] = {"enable_projects_ui"} _EFFECTIVE_UI_SETTINGS_CLASS: type[UISettings] | None = None -def register_extra_ui_setting(name: str, annotation: Any, field: FieldInfo) -> None: +def register_extra_ui_setting(name: str, annotation: object, field: FieldInfo) -> None: """Register an additional UI settings field contributed by an extension package. ``field`` must be a ``FieldInfo`` instance — construct it directly @@ -470,7 +537,7 @@ async def delete_allowed_ip( async def _get_settings_with_schema( settings_key: str, - settings_class: Any, + settings_class: type[BaseModel], config: dict, ) -> dict: """ @@ -842,7 +909,9 @@ async def get_sso_settings(): # Resolve the effective SSO config: the stored row wins, else the process # environment, else each field's default. Unlike the legacy read path this # does not write os.environ; a GET has no business mutating the environment. - sso_db_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) + sso_db_record: Final = await _sso_settings_mapping_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) sso_db_settings: Final = dict(sso_db_record.sso_settings) if sso_db_record and sso_db_record.sso_settings else None resolved: Final = resolve_sso_config(sso_db_settings, os.environ) @@ -914,8 +983,10 @@ async def update_sso_settings( # before-snapshot has the same shape as after_value, and rely on # create_config_audit_log's secret-name redaction to mask the # *_client_secret fields before the audit row is written. - existing_sso_record: Final = await SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}) - before_sso_data: dict[str, Any] | None = None + existing_sso_record: Final = await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).find_unique( + where={"id": "sso_config"} + ) + before_sso_data: dict[str, JsonValue] | None = None if existing_sso_record and existing_sso_record.sso_settings: stored = existing_sso_record.sso_settings if isinstance(stored, str): @@ -948,7 +1019,7 @@ async def update_sso_settings( encrypted_sso_data: Final = proxy_config._encrypt_env_variables(environment_variables=sso_data) # Save to dedicated SSO table - await SSOConfigRepository(prisma_client).table.upsert( + await _stored_sso_settings_db(SSOConfigRepository(prisma_client)).upsert( where={"id": "sso_config"}, data={ "create": { @@ -974,7 +1045,7 @@ async def update_sso_settings( # Remove SSO-related env vars from config.environment_variables try: - env_var_entry: Final = await ConfigRepository(prisma_client).table.find_unique( + env_var_entry: Final = await _config_param_db(ConfigRepository(prisma_client)).find_unique( where={"param_name": "environment_variables"} ) @@ -982,7 +1053,7 @@ async def update_sso_settings( if env_var_entry is not None: if env_var_entry.param_value is not None: if isinstance(env_var_entry.param_value, str): - environment_variables = json.loads(env_var_entry.param_value) + environment_variables: Mapping[str, object] = json.loads(env_var_entry.param_value) else: environment_variables = dict(env_var_entry.param_value) else: @@ -993,7 +1064,7 @@ async def update_sso_settings( key: value for key, value in environment_variables.items() if key not in env_vars_to_remove } - await ConfigRepository(prisma_client).table.update( + await _config_param_db(ConfigRepository(prisma_client)).update( where={"param_name": "environment_variables"}, data={ "param_value": json.dumps(filtered_env_vars, default=str), @@ -1239,8 +1310,10 @@ async def get_ui_settings_cached() -> dict[str, Any]: if prisma_client is None: return {} - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) - ui_settings: dict[str, Any] = {} + db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) + ui_settings: dict[str, JsonValue] = {} if db_record and db_record.ui_settings: raw: Final = db_record.ui_settings ui_settings = json.loads(raw) if isinstance(raw, str) else dict(raw) @@ -1272,9 +1345,11 @@ async def get_ui_settings(): detail={"error": "Database not connected. Please connect a database."}, ) - ui_settings: dict[str, Any] = {} + ui_settings: Mapping[str, JsonValue] = {} - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + 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 @@ -1300,7 +1375,7 @@ async def get_ui_settings(): await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) # Build config-like object for schema helper - config: Final[dict[str, Any]] = {"litellm_settings": {"ui_settings": ui_settings}} + config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}} return await _get_settings_with_schema( settings_key="ui_settings", @@ -1315,7 +1390,7 @@ async def get_ui_settings(): dependencies=[Depends(user_api_key_auth)], ) async def update_ui_settings( - settings_body: dict[str, Any] = Body(...), + settings_body: dict[str, object] = Body(...), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -1352,7 +1427,7 @@ async def update_ui_settings( raise HTTPException(status_code=422, detail=e.errors()) # Only include fields the caller actually sent (not Pydantic defaults). - settings_dict: Final = settings.model_dump(exclude_unset=True) + settings_dict: Final[Mapping[str, JsonValue]] = settings.model_dump(exclude_unset=True) # Reject enterprise-only settings up front so the caller gets a clear # signal instead of a silent drop. @@ -1373,15 +1448,17 @@ async def update_ui_settings( # Merge with existing persisted settings so a partial PATCH doesn't # overwrite fields the caller didn't send. - existing: dict = {} - db_existing: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + existing: dict[str, JsonValue] = {} + db_existing: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( + where={"id": "ui_settings"} + ) if db_existing and db_existing.ui_settings: raw: Final = db_existing.ui_settings existing = json.loads(raw) if isinstance(raw, str) else dict(raw) ui_settings: Final = {**existing, **incoming} - await UISettingsRepository(prisma_client).table.upsert( + await _ui_settings_db(UISettingsRepository(prisma_client)).upsert( where={"id": "ui_settings"}, data={ "create": { diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cdd41d2ed42..e59c6adaf22 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -165,6 +165,7 @@ if TYPE_CHECKING: from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -3000,6 +3001,13 @@ class PrismaClient: ] = [] # mutable-ok: drained queue, mirrors tool_usage_transactions _autorouter_turn_transactions_lock = asyncio.Lock() + # How long a health probe failure waits for an in-flight planned engine + # replacement to settle before deciding whether to report itself. Generous + # against a replacement that takes well under a second, and far short of the + # reconnect budget an outage-hung `connect()` runs under, so a real outage + # is never waited out. + PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS: ClassVar[float] = 5.0 + def __init__( self, database_url: str, @@ -4970,6 +4978,101 @@ class PrismaClient: else: verbose_proxy_logger.debug("Prisma DB health watchdog observed non-DB error: %s", e) + def _probe_target_wrapper(self) -> PrismaWrapper: + """The Prisma wrapper a `SELECT 1` health probe actually reaches. + + `health_check()` issues `query_raw`, which `RoutingPrismaWrapper` sends + to the reader unless the reader is degraded. The writer's engine state + therefore says nothing about a probe that failed against the reader, so + the gate has to follow the same routing rule the probe did. + """ + if isinstance(self.db, RoutingPrismaWrapper): + return self.db.writer if self.db.reader_unavailable else self.db.reader + return self.db + + async def _run_health_probe(self, wrapper: PrismaWrapper) -> object: + """Issue the `SELECT 1` a health check is made of, against `wrapper`. + + Takes the wrapper rather than re-reading `self.db`, because routing is + re-resolved on every attribute access: a reader that recovers between + the caller picking its target and the query going out would send the + probe to a different engine than the one whose generation the caller is + about to check, and attribute the failure to the wrong replacement. + """ + sql_query: Final = "SELECT 1" + response: Final = await wrapper.query_raw(sql_query) + return response + + async def _probe_answers_now(self, wrapper: PrismaWrapper) -> bool: + try: + await self._run_health_probe(wrapper) + except Exception as probe_error: # noqa: BLE001 # any failure means the database is not answering + verbose_proxy_logger.debug("Prisma health_check() confirmation probe failed: %s", probe_error) + return False + return True + + async def _planned_engine_replacement_absorbed( + self, + e: Exception, + wrapper: PrismaWrapper, + generation_before: int, + ) -> bool: + """True iff `e` is a connection-class probe failure that a completed + planned query-engine replacement explains. + + Planned replacements (RDS IAM token refresh, guarded reconnect) kill the + running query engine and spawn a new one. A `SELECT 1` probe that races + that sub-second window fails with a transport error against the engine's + local HTTP port even though nothing is wrong with the database, and + reporting it drives a false-positive `db_exceptions` alert on every + replacement. + + Two things must both hold, because neither is sufficient alone. The + engine generation must have moved, which says a replacement completed + rather than merely being attempted: reconnect attempts during a real + outage hold the same lock for tens of seconds, so gating on an in-flight + replacement would swallow most of an outage's alerts. And a fresh probe + must succeed, because `Prisma.connect()` polls the query engine's own + `/status` endpoint rather than round-tripping to the database, so a + future engine that binds before it validates its connection pool would + let the generation advance with the database still unreachable. + + Waiting for an in-flight replacement to settle is what makes the + generation check meaningful, since the generation has not moved yet at + the instant the probe fails. The wait is generous against a replacement + that takes well under a second and short enough that an outage-hung + reconnect is not waited out; a replacement that has not settled by then + reports rather than stays silent. + """ + if not PrismaDBExceptionHandler.is_database_connection_error(e): + return False + await wrapper.wait_for_planned_engine_replacement(self.PLANNED_ENGINE_REPLACEMENT_SETTLE_SECONDS) + if wrapper.engine_generation == generation_before: + return False + return await self._probe_answers_now(wrapper) + + async def _report_health_check_failure( + self, + e: Exception, + duration: float, + traceback_str: str, + wrapper: PrismaWrapper, + generation_before: int, + ) -> None: + if await self._planned_engine_replacement_absorbed(e, wrapper, generation_before): + verbose_proxy_logger.info( + "Prisma health_check() connection error raced a planned query-engine replacement; " + "not reporting it as a DB exception: %s", + e, + ) + return + await self.proxy_logging_obj.failure_handler( + original_exception=e, + duration=duration, + call_type="health_check", + traceback_str=traceback_str, + ) + @backoff.on_exception( backoff.expo, Exception, @@ -4982,13 +5085,10 @@ class PrismaClient: Health check endpoint for the prisma client """ start_time: Final = time.time() + probe_wrapper: Final = self._probe_target_wrapper() + generation_before: Final = probe_wrapper.engine_generation try: - sql_query: Final = "SELECT 1" - - # Execute the raw query - # The asterisk before `user_id_list` unpacks the list into separate arguments - response: Final = await self.db.query_raw(sql_query) - return response + return await self._run_health_probe(probe_wrapper) except Exception as e: import traceback @@ -4998,11 +5098,12 @@ class PrismaClient: end_time: Final = time.time() _duration: Final = end_time - start_time asyncio.create_task( - self.proxy_logging_obj.failure_handler( - original_exception=e, + self._report_health_check_failure( + e=e, duration=_duration, - call_type="health_check", traceback_str=error_traceback, + wrapper=probe_wrapper, + generation_before=generation_before, ) ) raise e @@ -6319,6 +6420,74 @@ def construct_database_url_from_env_vars() -> str | None: return None +async def _get_validated_team_object( + user_api_key_dict: "UserAPIKeyAuth", + team_id: str, + prisma_client: "PrismaClient", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging", +) -> "LiteLLM_TeamTableCachedObj": + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.management_endpoints.team_endpoints import validate_membership + + team_object: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) + return team_object + + +async def _get_team_object_for_access_groups( + team_id: str | None, + prisma_client: Optional["PrismaClient"], + user_api_key_cache: Optional["UserApiKeyCache"], + proxy_logging_obj: Optional["ProxyLogging"], +) -> Optional["LiteLLM_TeamTableCachedObj"]: + from litellm.proxy.auth.auth_checks import get_team_object + + if team_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None: + return None + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + verbose_proxy_logger.debug("Could not fetch team %s while listing models", team_id) + return None + + +async def _get_access_group_models( + user_api_key_dict: "UserAPIKeyAuth", + team_object: Optional["LiteLLM_TeamTableCachedObj"], + prisma_client: Optional["PrismaClient"], + user_api_key_cache: Optional["UserApiKeyCache"], + proxy_logging_obj: Optional["ProxyLogging"], +) -> tuple[str, ...]: + from litellm.proxy.auth.auth_checks import ( + _get_models_from_access_groups, + get_authorized_resources_from_key_access_groups, + ) + + team_group_models: Final = await _get_models_from_access_groups( + access_group_ids=(team_object.access_group_ids or ()) if team_object is not None else (), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + key_group_models: Final = await get_authorized_resources_from_key_access_groups( + valid_token=user_api_key_dict, + team_object=team_object, + resource_field="access_model_names", + ) + return tuple(dict.fromkeys((*team_group_models, *key_group_models))) + + async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], @@ -6350,13 +6519,11 @@ async def get_available_models_for_user( Returns: List of model names available to the user """ - from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.model_checks import ( get_complete_model_list, get_key_models, get_team_models, ) - from litellm.proxy.management_endpoints.team_endpoints import validate_membership # Get proxy model list and access groups if llm_router is None: @@ -6366,31 +6533,33 @@ async def get_available_models_for_user( proxy_model_list = llm_router.get_model_names() model_access_groups = llm_router.get_model_access_groups() - # Get key models - key_models = get_key_models( - user_api_key_dict=user_api_key_dict, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - ) - - # Get team models - team_models: list[str] = user_api_key_dict.team_models - - # If specific team_id is provided, validate and get team models - if team_id and prisma_client and proxy_logging_obj and user_api_key_cache: - key_models = [] - team_object: Final = await get_team_object( + requested_team_object: Final = ( + await _get_validated_team_object( + user_api_key_dict=user_api_key_dict, team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) - team_models = team_object.models + if team_id and prisma_client and proxy_logging_obj and user_api_key_cache + else None + ) - team_models = get_team_models( - team_models=team_models, + key_models: Final[Sequence[str]] = ( + () + if requested_team_object is not None + else get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + ) + ) + + team_models: Final = get_team_models( + team_models=( + requested_team_object.models if requested_team_object is not None else user_api_key_dict.team_models + ), proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, include_model_access_groups=include_model_access_groups, @@ -6398,10 +6567,31 @@ async def get_available_models_for_user( effective_team_id: Final = team_id or user_api_key_dict.team_id + access_group_models: Final = ( + await _get_access_group_models( + user_api_key_dict=user_api_key_dict, + team_object=requested_team_object + or await _get_team_object_for_access_groups( + team_id=effective_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if key_models or team_models + else () + ) + + granted_key_models: Final = (*key_models, *access_group_models) if key_models else key_models + granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models + # Get complete model list all_models: Final = get_complete_model_list( - key_models=key_models, - team_models=team_models, + key_models=granted_key_models, + team_models=granted_team_models, proxy_model_list=proxy_model_list, user_model=user_model, infer_model_from_keys=general_settings.get("infer_model_from_keys", False), diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index 6b40281198c..2b037bef795 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -10,10 +10,17 @@ All /vector_store management endpoints import copy import json -from typing import Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import APIRouter, Depends, HTTPException +if TYPE_CHECKING: + from prisma.models import LiteLLM_ManagedVectorStoresTable as _VectorStoreRow + + from litellm.proxy.utils import PrismaClient + from litellm.router import Router + import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -43,6 +50,25 @@ from litellm.vector_stores.vector_store_registry import VectorStoreRegistry router: Final = APIRouter() + +class _VectorStoreTableActions(Protocol): + async def find_unique(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ... + + async def create(self, data: Mapping[str, object]) -> "_VectorStoreRow": ... + + async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> "_VectorStoreRow": ... + + async def delete(self, where: Mapping[str, str]) -> "_VectorStoreRow | None": ... + + +def _vector_store_table(prisma_client: "PrismaClient") -> _VectorStoreTableActions: + return ManagedVectorStoresRepository(prisma_client).table + + +def _row_to_vector_store(row: "_VectorStoreRow") -> LiteLLM_ManagedVectorStore: + return LiteLLM_ManagedVectorStore(**row.model_dump()) + + _LITELLM_PARAMS_MASKER: Final = SensitiveDataMasker() @@ -117,22 +143,20 @@ def _redact_sensitive_litellm_params(litellm_params: Any, _depth: int = 0) -> An async def _fetch_and_authorize_vector_store( vector_store_id: str, user_api_key_dict: UserAPIKeyAuth, - prisma_client: Any, + prisma_client: "PrismaClient", ) -> "LiteLLM_ManagedVectorStore": """ Look up a vector store by id and confirm the caller can access it. Raises HTTPException(404) on miss and HTTPException(403) on access denial. """ - row: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( - where={"vector_store_id": vector_store_id} - ) + row: Final = await _vector_store_table(prisma_client).find_unique(where={"vector_store_id": vector_store_id}) if row is None: raise HTTPException( status_code=404, detail=f"Vector store with ID {vector_store_id} not found", ) - typed: Final = LiteLLM_ManagedVectorStore(**row.model_dump()) + typed: Final = _row_to_vector_store(row) if not await _check_vector_store_access(typed, user_api_key_dict): raise HTTPException( status_code=403, @@ -141,7 +165,7 @@ async def _fetch_and_authorize_vector_store( return typed -def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, Any] | None: +def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> dict[str, object] | None: """ Resolve embedding config from router's config-defined models. @@ -177,7 +201,7 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d litellm_params = deployment.litellm_params # Build embedding config from model params - embedding_config: dict[str, Any] = {} + embedding_config: dict[str, object] = {} # Extract api_key api_key = getattr(litellm_params, "api_key", None) @@ -217,7 +241,9 @@ def _resolve_embedding_config_from_router(embedding_model: str, llm_router) -> d return None -async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) -> dict[str, Any] | None: +async def _resolve_embedding_config_from_db( + embedding_model: str, prisma_client: "PrismaClient" +) -> dict[str, object] | None: """ Resolve embedding config from database model configuration. @@ -307,7 +333,9 @@ async def _resolve_embedding_config_from_db(embedding_model: str, prisma_client) return None -async def _resolve_embedding_config(embedding_model: str, prisma_client, llm_router=None) -> dict[str, Any] | None: +async def _resolve_embedding_config( + embedding_model: str, prisma_client: "PrismaClient | None", llm_router: "Router | None" = None +) -> dict[str, object] | None: """ Resolve embedding config from either router (config-defined) or database models. @@ -388,7 +416,7 @@ async def _check_vector_store_access( async def create_vector_store_in_db( vector_store_id: str, custom_llm_provider: str, - prisma_client, + prisma_client: "PrismaClient | None", vector_store_name: str | None = None, vector_store_description: str | None = None, vector_store_metadata: dict | None = None, @@ -417,7 +445,7 @@ async def create_vector_store_in_db( raise HTTPException(status_code=500, detail="Database not connected") # Check if vector store already exists - existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique( where={"vector_store_id": vector_store_id} ) if existing_vector_store is not None: @@ -427,7 +455,7 @@ async def create_vector_store_in_db( ) # Prepare data for database - data_to_create: Final[dict[str, Any]] = { + data_to_create: Final[dict[str, object]] = { "vector_store_id": vector_store_id, "custom_llm_provider": custom_llm_provider, } @@ -463,9 +491,9 @@ async def create_vector_store_in_db( data_to_create["litellm_params"] = safe_dumps({}) # Create in database - _new_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.create(data=data_to_create) + _new_vector_store: Final = await _vector_store_table(prisma_client).create(data=data_to_create) - new_vector_store: Final[LiteLLM_ManagedVectorStore] = LiteLLM_ManagedVectorStore(**_new_vector_store.model_dump()) + new_vector_store: Final[LiteLLM_ManagedVectorStore] = _row_to_vector_store(_new_vector_store) # Add vector store to registry if litellm.vector_store_registry is not None: @@ -682,12 +710,12 @@ async def delete_vector_store( memory_vector_store_exists = False vector_store_to_check = None - existing_vector_store: Final = await ManagedVectorStoresRepository(prisma_client).table.find_unique( + existing_vector_store: Final = await _vector_store_table(prisma_client).find_unique( where={"vector_store_id": data.vector_store_id} ) if existing_vector_store is not None: db_vector_store_exists = True - vector_store_to_check = LiteLLM_ManagedVectorStore(**existing_vector_store.model_dump()) + vector_store_to_check = _row_to_vector_store(existing_vector_store) # Check in-memory registry if litellm.vector_store_registry is not None: @@ -715,9 +743,7 @@ async def delete_vector_store( # Delete from database if exists if db_vector_store_exists: - await ManagedVectorStoresRepository(prisma_client).table.delete( - where={"vector_store_id": data.vector_store_id} - ) + await _vector_store_table(prisma_client).delete(where={"vector_store_id": data.vector_store_id}) # Delete from in-memory registry if exists if memory_vector_store_exists and litellm.vector_store_registry is not None: @@ -829,7 +855,7 @@ async def update_vector_store( try: update_data: Final = data.model_dump(exclude_unset=True) - vector_store_id: Final = update_data.pop("vector_store_id") + vector_store_id: Final[str] = update_data.pop("vector_store_id") # Per-store access control: anyone authenticated who passes the # premium-feature gate could otherwise update *any* vector store — @@ -857,12 +883,12 @@ async def update_vector_store( update_data["litellm_params"] = safe_dumps(litellm_params_dict) # Update in database - updated: Final = await ManagedVectorStoresRepository(prisma_client).table.update( + updated: Final = await _vector_store_table(prisma_client).update( where={"vector_store_id": vector_store_id}, data=update_data, ) - updated_vs: Final = LiteLLM_ManagedVectorStore(**updated.model_dump()) + updated_vs: Final = _row_to_vector_store(updated) # Immediately update in-memory registry to keep it in sync if litellm.vector_store_registry is not None: diff --git a/litellm/proxy/vector_store_files_endpoints/endpoints.py b/litellm/proxy/vector_store_files_endpoints/endpoints.py index 896b7ca33d7..c9b89bcd390 100644 --- a/litellm/proxy/vector_store_files_endpoints/endpoints.py +++ b/litellm/proxy/vector_store_files_endpoints/endpoints.py @@ -30,10 +30,12 @@ if TYPE_CHECKING: router: Final = APIRouter() -def _update_request_data_with_managed_file_id( +async def _update_request_data_with_managed_file_id( data: dict, file_id: str, request: Request, + user_api_key_dict: UserAPIKeyAuth, + managed_files_obj: object | None, llm_router: Optional["Router"] = None, ) -> tuple[dict, str | None]: """ @@ -65,6 +67,16 @@ def _update_request_data_with_managed_file_id( is_base64_encoded_unified_id, parse_unified_id, ) + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_id_requirement, + ) + + await validate_managed_id_requirement( + resource_id=file_id, + resource_kind="file", + user_api_key_dict=user_api_key_dict, + managed_files_obj=managed_files_obj, + ) # First, check if this is a unified managed file ID (base64 encoded) decoded_id: Final = is_base64_encoded_unified_id(file_id) @@ -509,8 +521,13 @@ async def vector_store_file_create( # Handle managed file IDs if present in request body original_managed_file_id = None if "file_id" in data: - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=data["file_id"], request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=data["file_id"], + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -707,8 +724,13 @@ async def vector_store_file_retrieve( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -809,8 +831,13 @@ async def vector_store_file_content( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -911,8 +938,13 @@ async def vector_store_file_update( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs @@ -1013,8 +1045,13 @@ async def vector_store_file_delete( ) # Handle managed file IDs first - data, original_managed_file_id = _update_request_data_with_managed_file_id( - data=data, file_id=file_id, request=request, llm_router=llm_router + data, original_managed_file_id = await _update_request_data_with_managed_file_id( + data=data, + file_id=file_id, + request=request, + user_api_key_dict=user_api_key_dict, + managed_files_obj=proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router=llm_router, ) # Then handle managed vector store IDs diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 66e0b6d59e7..be19f290ba6 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -42,6 +42,10 @@ class AgentsRepository(PrismaTableRepository): table_name = "litellm_agentstable" +class ObjectPermissionRepository(PrismaTableRepository): + table_name = "litellm_objectpermissiontable" + + class GuardrailsRepository(PrismaTableRepository): table_name = "litellm_guardrailstable" diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f54023836e5..b2d065ea23b 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -4,8 +4,8 @@ Handles transforming from Responses API -> LiteLLM completion (Chat Completion import json import re -from collections.abc import Sequence -from typing import Any, Final, Literal, cast +from collections.abc import Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable from openai.types.chat.chat_completion_named_tool_choice_param import ( ChatCompletionNamedToolChoiceParam, @@ -16,6 +16,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( from openai.types.responses import ResponseFunctionToolCall from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_param import FunctionToolParam +from pydantic import TypeAdapter from typing_extensions import TypedDict from litellm._logging import verbose_logger @@ -78,9 +79,35 @@ from .custom_tools import ( unwrap_custom_tool_arguments, ) +if TYPE_CHECKING: + from openai.types.responses.response_apply_patch_tool_call import ( + ResponseApplyPatchToolCall, + ) + ########### Initialize Classes used for Responses API ########### TOOL_CALLS_CACHE: Final = InMemoryCache() +_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) +_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) +_TEXT_ADAPTER: Final = TypeAdapter(str) + + +@runtime_checkable +class _SupportsIter(Protocol): + def __iter__(self) -> Iterator[object]: ... + + +@runtime_checkable +class _HasToolCalls(Protocol): + tool_calls: object + + +@runtime_checkable +class _HasId(Protocol): + id: object + class ChatCompletionSession(TypedDict, total=False): messages: list[ @@ -205,7 +232,7 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams, custom_llm_provider: str | None = None, stream: bool | None = None, - extra_headers: dict[str, Any] | None = None, + extra_headers: Mapping[str, object] | None = None, **kwargs, ) -> dict: """ @@ -462,7 +489,9 @@ class LiteLLMCompletionResponsesConfig: if not chat_completion_messages: continue - deduped_in_place: list[Any] = [] + deduped_in_place: list[ + AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage + ] = [] for m in chat_completion_messages: role = "" if isinstance(m, dict): @@ -472,7 +501,7 @@ class LiteLLMCompletionResponsesConfig: # Drop assistant tool_calls wrappers if we already have this call_id if role == "assistant": - tool_calls: Any = ( + tool_calls: object = ( m.get("tool_calls") if isinstance(m, dict) else getattr(m, "tool_calls", None) ) call_id = "" @@ -534,7 +563,7 @@ class LiteLLMCompletionResponsesConfig: call_id = "" if role == "assistant": - tool_calls: Any = None + tool_calls: object = None if isinstance(tool_call_message, dict): tool_calls = tool_call_message.get("tool_calls") else: @@ -578,7 +607,16 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _find_previous_assistant_idx(messages: list[Any], current_idx: int) -> int | None: + def _find_previous_assistant_idx( + messages: Sequence[ + AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message + ], + current_idx: int, + ) -> int | None: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": @@ -586,7 +624,18 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _recover_tool_call_id_from_assistant(assistant_message: Any, message: Any) -> str: + def _recover_tool_call_id_from_assistant( + assistant_message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + ) -> str: """Try to recover empty tool_call_id from assistant message's tool_calls.""" tool_calls_raw: Final = ( assistant_message.get("tool_calls") @@ -594,17 +643,23 @@ class LiteLLMCompletionResponsesConfig: else getattr(assistant_message, "tool_calls", None) ) if tool_calls_raw and isinstance(tool_calls_raw, list) and len(tool_calls_raw) > 0: - first_tool_call: Final = tool_calls_raw[0] + first_tool_call: Final = _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw)[0] if isinstance(first_tool_call, dict): - tool_call_id_raw = first_tool_call.get("id", "") + tool_call_id_raw = _ANY_KEY_DICT_ADAPTER.validate_python(first_tool_call).get("id", "") return str(tool_call_id_raw) if tool_call_id_raw is not None else "" - elif hasattr(first_tool_call, "id"): - tool_call_id_raw = getattr(first_tool_call, "id", None) + elif isinstance(first_tool_call, _HasId): + tool_call_id_raw = first_tool_call.id return str(tool_call_id_raw) if tool_call_id_raw is not None else "" return "" @staticmethod - def _get_tool_calls_list(assistant_message: Any) -> list[Any]: + def _get_tool_calls_list( + assistant_message: AllMessageValues + | GenericChatCompletionMessage + | ChatCompletionResponseMessage + | ChatCompletionMessageToolCall + | Message, + ) -> Sequence[object]: """Extract tool_calls as a list from assistant message.""" tool_calls_raw: Final = ( assistant_message.get("tool_calls") @@ -614,18 +669,18 @@ class LiteLLMCompletionResponsesConfig: if tool_calls_raw is None: return [] if isinstance(tool_calls_raw, list): - return tool_calls_raw - if hasattr(tool_calls_raw, "__iter__") and not isinstance(tool_calls_raw, (str, bytes)): + return _OBJECT_LIST_ADAPTER.validate_python(tool_calls_raw) + if isinstance(tool_calls_raw, _SupportsIter) and not isinstance(tool_calls_raw, (str, bytes)): return list(tool_calls_raw) return [] @staticmethod - def _check_tool_call_exists(tool_calls: list[Any], tool_call_id: str) -> bool: + def _check_tool_call_exists(tool_calls: Sequence[object], tool_call_id: str) -> bool: """Check if a tool_call with the given ID exists in the list.""" for tool_call in tool_calls: - tool_call_id_to_check: str | None = None + tool_call_id_to_check: object = None if isinstance(tool_call, dict): - tool_call_id_to_check = tool_call.get("id") + tool_call_id_to_check = _ANY_KEY_DICT_ADAPTER.validate_python(tool_call).get("id") elif hasattr(tool_call, "id"): tool_call_id_to_check = getattr(tool_call, "id", None) if tool_call_id_to_check == tool_call_id: @@ -633,12 +688,13 @@ class LiteLLMCompletionResponsesConfig: return False @staticmethod - def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: list[Any]) -> dict[str, Any] | None: + def _reconstruct_tool_call_from_tools(tool_call_id: str, tools: Sequence[object]) -> dict[str, object] | None: """Reconstruct a minimal tool_call definition from tools list.""" for tool in tools: if isinstance(tool, dict): - tool_function = tool.get("function") or {} - tool_name = tool_function.get("name") or tool.get("name") or "" + tool_map = _ANY_KEY_DICT_ADAPTER.validate_python(tool) + tool_function = _ANY_KEY_DICT_ADAPTER.validate_python(tool_map.get("function") or {}) + tool_name = tool_function.get("name") or tool_map.get("name") or "" if tool_name: return { "id": tool_call_id, @@ -651,7 +707,7 @@ class LiteLLMCompletionResponsesConfig: return None @staticmethod - def _get_mapping_or_attr_value(obj: Any, key: str, default: Any = None) -> Any: + def _get_mapping_or_attr_value(obj: object, key: str, default: object = None) -> object: """ Safely read a field from dict-like or attribute-based objects. """ @@ -659,7 +715,7 @@ class LiteLLMCompletionResponsesConfig: return default if isinstance(obj, dict): - return obj.get(key, default) + return _ANY_KEY_DICT_ADAPTER.validate_python(obj).get(key, default) getter: Final = getattr(obj, "get", None) if callable(getter): @@ -672,13 +728,13 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _create_tool_call_chunk( - tool_use_definition: dict[str, Any], tool_call_id: str, index: int + tool_use_definition: Mapping[object, object], tool_call_id: str, index: int ) -> ChatCompletionToolCallChunk: """Create a ChatCompletionToolCallChunk from tool_use_definition.""" function_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "function") function_name_raw: Final = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "name") function_arguments_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(function_raw, "arguments") - function: Final[dict[str, Any]] = { + function: Final[dict[str, object]] = { "name": function_name_raw or "", "arguments": function_arguments_raw or "{}", } @@ -697,7 +753,7 @@ class LiteLLMCompletionResponsesConfig: ) @staticmethod - def _normalize_tool_use_definition(tool_use_definition: Any, tool_call_id: str) -> dict[str, Any] | None: + def _normalize_tool_use_definition(tool_use_definition: object, tool_call_id: str) -> dict[object, object] | None: """ Normalize cached tool_call definitions to a dict-like shape consumed by _create_tool_call_chunk. """ @@ -705,7 +761,7 @@ class LiteLLMCompletionResponsesConfig: return None if isinstance(tool_use_definition, dict): - normalized_definition: dict[str, Any] = dict(tool_use_definition) + normalized_definition: dict[object, object] = _ANY_KEY_DICT_ADAPTER.validate_python(tool_use_definition) else: tool_use_id_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "id") tool_use_type_raw = LiteLLMCompletionResponsesConfig._get_mapping_or_attr_value(tool_use_definition, "type") @@ -738,7 +794,7 @@ class LiteLLMCompletionResponsesConfig: return normalized_definition @staticmethod - def _add_tool_call_to_assistant(assistant_message: Any, tool_call_chunk: ChatCompletionToolCallChunk) -> None: + def _add_tool_call_to_assistant(assistant_message: object, tool_call_chunk: ChatCompletionToolCallChunk) -> None: """Add a tool_call to an assistant message.""" if isinstance(assistant_message, dict): prev_assistant_dict: Final = cast(dict[str, Any], assistant_message) @@ -747,7 +803,7 @@ class LiteLLMCompletionResponsesConfig: tool_calls_list: Final = prev_assistant_dict["tool_calls"] if isinstance(tool_calls_list, list): tool_calls_list.append(tool_call_chunk) - elif hasattr(assistant_message, "tool_calls"): + elif isinstance(assistant_message, _HasToolCalls): if assistant_message.tool_calls is None: assistant_message.tool_calls = [] if isinstance(assistant_message.tool_calls, list): @@ -762,7 +818,7 @@ class LiteLLMCompletionResponsesConfig: | ChatCompletionMessageToolCall | Message ], - tools: list[Any] | None = None, + tools: Sequence[object] | None = None, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -851,7 +907,7 @@ class LiteLLMCompletionResponsesConfig: tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list(prev_assistant) if not LiteLLMCompletionResponsesConfig._check_tool_call_exists(tool_calls, tool_call_id): - _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) + _tool_use_definition: object = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) if not _tool_use_definition and tools: _tool_use_definition = LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools( @@ -908,7 +964,7 @@ class LiteLLMCompletionResponsesConfig: function_call=input_item ) else: - content: Final = input_item.get("content") + content: Final[object] = input_item.get("content") # Handle None content: Responses API allows None content, but GenericChatCompletionMessage requires content # Since guardrails skip None content anyway, we return empty list to exclude it from structured messages if content is None: @@ -923,7 +979,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_tool_call_output(input_item: Any) -> bool: + def _is_input_item_tool_call_output(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a tool call output """ @@ -936,7 +992,7 @@ class LiteLLMCompletionResponsesConfig: ] @staticmethod - def _is_input_item_function_call(input_item: Any) -> bool: + def _is_input_item_function_call(input_item: Mapping[str, object]) -> bool: """ Check if the input item is a function call or custom tool call. Both need to be reconstructed as assistant tool_calls for Chat @@ -946,7 +1002,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_tool_call_output_to_chat_completion_message( - tool_call_output: dict[str, Any], + tool_call_output: Mapping[str, object], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ ChatCompletionToolMessage is used to indicate the output from a tool call @@ -958,7 +1014,7 @@ class LiteLLMCompletionResponsesConfig: return [] def _normalize_function_call_output_to_tool_content( - output: Any, + output: object, ) -> Any: """ Normalize Responses API function_call_output.output into a shape that downstream @@ -981,7 +1037,7 @@ class LiteLLMCompletionResponsesConfig: # Some adapters represent tool output as a list of "input_*" parts if isinstance(output, list): - normalized_blocks: Final[list[dict[str, Any]]] = [] + normalized_blocks: Final[list[dict[str, object]]] = [] text_acc: Final[list[str]] = [] for part in output: if not isinstance(part, dict): @@ -1082,7 +1138,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_function_call_to_chat_completion_message( - function_call: dict[str, Any], + function_call: Mapping[str, str], ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API function_call into a Chat Completion message with tool calls @@ -1127,7 +1183,7 @@ class LiteLLMCompletionResponsesConfig: return [chat_completion_response_message] @staticmethod - def _resolve_file_id(item: dict[str, Any]) -> str | None: + def _resolve_file_id(item: Mapping[str, object]) -> object: """ Return the effective file_id for a Responses API input_file item. Explicit file_id takes precedence; file_url is used as fallback so @@ -1136,7 +1192,7 @@ class LiteLLMCompletionResponsesConfig: return item.get("file_id") or item.get("file_url") or None @staticmethod - def _transform_input_file_item_to_file_item(item: dict[str, Any]) -> dict[str, Any]: + def _transform_input_file_item_to_file_item(item: Mapping[str, object]) -> dict[str, object]: """ Transform a Responses API input_file item to a Chat Completion file item @@ -1146,21 +1202,21 @@ class LiteLLMCompletionResponsesConfig: Returns: Dictionary with transformed file structure for Chat Completion """ - file_dict: Final[dict[str, Any]] = {} + file_dict: Final[dict[str, object]] = {} file_id: Final = LiteLLMCompletionResponsesConfig._resolve_file_id(item) if file_id: file_dict["file_id"] = file_id if item.get("file_data"): file_dict["file_data"] = item["file_data"] - new_item: Final[dict[str, Any]] = {"type": "file", "file": file_dict} + new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: new_item["cache_control"] = item["cache_control"] return new_item @staticmethod def _transform_input_image_item_to_image_item( - item: dict[str, Any], + item: Mapping[str, str], ) -> ChatCompletionImageObject: """ Transform a Responses API input_image item to a Chat Completion image item @@ -1173,8 +1229,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_responses_api_content_to_chat_completion_content( - content: Any, - ) -> str | list[str | dict[str, Any]]: + content: object, + ) -> str | list[str | dict[str, object]]: """ Transform a Responses API content into a Chat Completion content @@ -1188,7 +1244,7 @@ class LiteLLMCompletionResponsesConfig: elif isinstance(content, str): return content elif isinstance(content, list): - content_list: Final[list[str | dict[str, Any]]] = [] + content_list: Final[list[str | dict[str, object]]] = [] for item in content: if isinstance(item, str): content_list.append(item) @@ -1198,8 +1254,8 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item(item) ) elif item.get("type") == "input_image": - image_block = dict( - LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item) + image_block = _STR_KEY_DICT_ADAPTER.validate_python( + dict(LiteLLMCompletionResponsesConfig._transform_input_image_item_to_image_item(item)) ) if "cache_control" in item: image_block["cache_control"] = item["cache_control"] @@ -1209,7 +1265,7 @@ class LiteLLMCompletionResponsesConfig: text_value = item.get("text") if text_value is None: continue - content_block: dict[str, Any] = { + content_block: dict[str, object] = { "type": LiteLLMCompletionResponsesConfig._get_chat_completion_request_content_type( item.get("type") or "text" ), @@ -1299,7 +1355,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(typed_tool.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - chat_completion_tool: dict[str, Any] = { + chat_completion_tool: dict[str, object] = { "type": "function", "function": { "name": typed_tool.get("name") or "", @@ -1340,7 +1396,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def transform_chat_completion_tool_params_to_responses_api_tools( chat_completion_tools: list[ChatCompletionToolParam | OpenAIMcpServerTool] | None, - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """ Transform Chat Completion tool params (e.g. from guardrail output) back to Responses API request tool format. Inverse of @@ -1348,7 +1404,7 @@ class LiteLLMCompletionResponsesConfig: """ if chat_completion_tools is None or not chat_completion_tools: return [] - result: Final[list[dict[str, Any]]] = [] + result: Final[list[dict[str, object]]] = [] for tool in chat_completion_tools: if not isinstance(tool, dict): result.append(tool) @@ -1358,7 +1414,7 @@ class LiteLLMCompletionResponsesConfig: parameters = dict(fn.get("parameters", {}) or {}) if not parameters or "type" not in parameters: parameters["type"] = "object" - responses_tool: dict[str, Any] = { + responses_tool: dict[str, object] = { "type": "function", "name": fn.get("name") or "", "description": fn.get("description") or "", @@ -1510,7 +1566,7 @@ class LiteLLMCompletionResponsesConfig: def convert_response_function_tool_call_to_chat_completion_tool_call( tool_call_item: Any, index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseFunctionToolCall to ChatCompletionToolCallChunk format. @@ -1536,7 +1592,7 @@ class LiteLLMCompletionResponsesConfig: else (dict(provider_fields) if hasattr(provider_fields, "__dict__") else {}) ) - function_dict: Final[dict[str, Any]] = { + function_dict: Final[dict[str, object]] = { "name": tool_call_item.name, "arguments": tool_call_item.arguments, } @@ -1544,7 +1600,7 @@ class LiteLLMCompletionResponsesConfig: if provider_specific_fields: function_dict["provider_specific_fields"] = provider_specific_fields - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": LiteLLMCompletionResponsesConfig._tool_call_id_from_responses_item( getattr(tool_call_item, "id", None), getattr(tool_call_item, "call_id", None), @@ -1561,9 +1617,9 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def convert_apply_patch_tool_call_to_chat_completion_tool_call( - tool_call_item: Any, + tool_call_item: "ResponseApplyPatchToolCall", index: int = 0, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. @@ -1581,7 +1637,7 @@ class LiteLLMCompletionResponsesConfig: import json operation_dict: Final = tool_call_item.operation.model_dump() - tool_call_dict: Final[dict[str, Any]] = { + tool_call_dict: Final[dict[str, object]] = { "id": tool_call_item.call_id, "function": { "name": "apply_patch", @@ -1795,9 +1851,11 @@ class LiteLLMCompletionResponsesConfig: if not images: return image_generation_items - for idx, image_item in enumerate(images): + for idx, image_item in enumerate(_DICT_ITEMS_LIST_ADAPTER.validate_python(images)): # Extract base64 from data URL - image_url = image_item.get("image_url", {}).get("url", "") + image_url = _TEXT_ADAPTER.validate_python( + _ANY_KEY_DICT_ADAPTER.validate_python(image_item.get("image_url", {})).get("url", "") + ) base64_data = LiteLLMCompletionResponsesConfig._extract_base64_from_data_url(image_url) if base64_data: @@ -2048,8 +2106,8 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_text_format_to_response_format( - text_param: dict[str, Any] | Any, - ) -> dict[str, Any] | None: + text_param: object, + ) -> dict[str, object] | None: """ Transform Responses API text.format parameter to Chat Completion response_format parameter. diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f923702119c..7b02c1b8023 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1064,6 +1064,7 @@ def responses( extra_headers=extra_headers, extra_body=extra_body, timeout=timeout if timeout is not None else request_timeout, + allowed_openai_params=allowed_openai_params, **kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 820839fc6bf..2e1e1a44594 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,7 +5,7 @@ import json import time import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -33,6 +33,7 @@ from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfi from litellm.responses.utils import ResponseAPILoggingUtils, ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PART_UNION_TYPES, + ResponseAPIUsage, ResponsesAPIResponse, ResponsesAPIStreamEvents, ResponsesAPIStreamingResponse, @@ -112,7 +113,7 @@ _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType( def _error_event_fields(error_obj: object) -> tuple[str, str | None, str | None]: - if isinstance(error_obj, dict): + if _is_json_object(error_obj): raw_message = error_obj.get("message") raw_type = error_obj.get("type") raw_code = error_obj.get("code") @@ -243,7 +244,9 @@ class BaseResponsesAPIStreamingIterator: # Using getattr(..., "response") alone is unsafe with Mocks: they synthesize a # truthy child Mock for any attribute, which breaks tests and is wrong on stream. if "response" in parsed_chunk: - response_object: Final = getattr(openai_responses_api_chunk, "response", None) + response_object: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_object is not None: response: Final = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id( responses_api_response=response_object, @@ -279,7 +282,9 @@ class BaseResponsesAPIStreamingIterator: model_id=_stream_model_id, ) elif _event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE: - _part: Final = getattr(openai_responses_api_chunk, "part", None) + _part: Final[PART_UNION_TYPES | Mapping[str, object] | None] = getattr( + openai_responses_api_chunk, "part", None + ) if _part is not None: if isinstance(_part, dict): ResponsesAPIRequestUtils._encode_container_ids_in_annotations( @@ -302,7 +307,7 @@ class BaseResponsesAPIStreamingIterator: openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, openai_types.ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, ): - item: Final = getattr(openai_responses_api_chunk, "item", None) + item: Final[object | None] = getattr(openai_responses_api_chunk, "item", None) if item: encrypted_content: Final = getattr(item, "encrypted_content", None) if encrypted_content and isinstance(encrypted_content, str): @@ -324,9 +329,11 @@ class BaseResponsesAPIStreamingIterator: self.completed_response = openai_responses_api_chunk # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and self.logging_obj is not None: - response_obj: Final[Any | None] = getattr(openai_responses_api_chunk, "response", None) + response_obj: Final[ResponsesAPIResponse | None] = getattr( + openai_responses_api_chunk, "response", None + ) if response_obj: - usage_obj: Final[Any | None] = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is not None: try: cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) @@ -414,7 +421,9 @@ class BaseResponsesAPIStreamingIterator: async_failure_handler / failure_handler so logging integrations correctly record the call as failed. """ - response_obj: Final = getattr(self.completed_response, "response", None) if self.completed_response else None + response_obj: Final[ResponsesAPIResponse | None] = ( + getattr(self.completed_response, "response", None) if self.completed_response else None + ) error_info: Final = getattr(response_obj, "error", None) if response_obj else None error_message, error_type, error_code = _error_event_fields(error_info) self._record_failed_response_usage(response_obj) @@ -429,7 +438,7 @@ class BaseResponsesAPIStreamingIterator: def _record_failed_response_usage(self, response_obj: ResponsesAPIResponse | None) -> None: if response_obj is None or self.logging_obj is None: return - usage_obj: Final = getattr(response_obj, "usage", None) + usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is None: return try: @@ -506,7 +515,7 @@ class BaseResponsesAPIStreamingIterator: return request_kwargs = getattr(caching_handler, "request_kwargs", None) - if not isinstance(request_kwargs, dict) or request_kwargs.get("stream") is not True: + if not _is_json_object(request_kwargs) or request_kwargs.get("stream") is not True: return request_kwargs = request_kwargs.copy() preset_cache_key = getattr(caching_handler, "preset_cache_key", None) @@ -606,7 +615,7 @@ class BaseResponsesAPIStreamingIterator: if self.completed_response is None: return - request_payload: Final[dict[str, Any]] = {} + request_payload: Final[dict[str, object]] = {} if isinstance(self.request_data, dict): request_payload.update(self.request_data) try: @@ -695,11 +704,15 @@ class BaseResponsesAPIStreamingIterator: pass -async def call_post_streaming_hooks_for_testing(iterator, chunk): +async def call_post_streaming_hooks_for_testing( + iterator: object, chunk: ResponsesAPIStreamingResponse +) -> ResponsesAPIStreamingResponse: """ Module-level helper for tests to ensure hooks can be invoked even if the iterator is wrapped. """ - hook_fn: Final = getattr(iterator, "_call_post_streaming_deployment_hook", None) + hook_fn: Final[Callable[[ResponsesAPIStreamingResponse], Awaitable[ResponsesAPIStreamingResponse]] | None] = ( + getattr(iterator, "_call_post_streaming_deployment_hook", None) + ) if hook_fn is None: return chunk return await hook_fn(chunk) @@ -1019,7 +1032,7 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def _dump_response_object(obj: Any) -> dict[str, Any]: if hasattr(obj, "model_dump"): return obj.model_dump() - if isinstance(obj, dict): + if _is_json_object(obj): return obj return {} @@ -1684,7 +1697,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + evt_obj: Final[Mapping[str, object]] = json.loads(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -1925,7 +1938,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, Any], + completed_event: dict[str, object], ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into @@ -2065,7 +2078,7 @@ class ManagedResponsesWebSocketHandler: Flat: {"type": "response.create", "input": [...], "model": "...", ...} """ nested: Final = msg_obj.get("response") - response_params: Final[dict[str, Any]] = ( + response_params: Final[dict[str, object]] = ( nested if _is_json_object(nested) and nested else {k: v for k, v in msg_obj.items() if k != "type"} ) return { @@ -2076,7 +2089,7 @@ class ManagedResponsesWebSocketHandler: def _apply_history( self, - call_kwargs: dict[str, Any], + call_kwargs: dict[str, object], previous_response_id: str | None, current_messages: list[dict[str, object]], prior_history: list[dict[str, object]], @@ -2129,7 +2142,7 @@ class ManagedResponsesWebSocketHandler: return False return event_provider == self._connection_provider - def _inject_credentials(self, call_kwargs: dict[str, Any], model: str | None = None) -> None: + def _inject_credentials(self, call_kwargs: dict[str, object], model: str | None = None) -> None: """Inject connection-level credentials and metadata into call_kwargs.""" if self.api_key is not None: call_kwargs["api_key"] = self.api_key diff --git a/litellm/router.py b/litellm/router.py index 1edb80da7ce..1a76eb5d59b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -109,6 +109,7 @@ from litellm.router_utils.common_utils import ( filter_team_based_models, filter_web_search_deployments, resolve_model_group_alias, + truncate_fallback_error_detail, ) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( @@ -342,6 +343,12 @@ def _replay_live_router_model_cost() -> None: set_live_deployment_replay(_replay_live_router_model_cost) +# Kwargs that log_retry must not copy into a retry breadcrumb. The breadcrumbs reach spend +# logs and logging callbacks, and these carry either the request payload or router-internal +# walk state rather than anything that identifies the failed attempt. +RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset(("messages", "original_function", "attempted_targets")) + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -6361,17 +6368,16 @@ class Router: return response except Exception as new_exception: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) - fallback_failure_exception_str = redact_string(str(new_exception)) + fallback_failure_exception_str = truncate_fallback_error_detail(redact_string(str(new_exception))) cooldown_info: Final = await _async_get_cooldown_deployments_with_debug_info( litellm_router_instance=self, parent_otel_span=parent_otel_span, ) verbose_router_logger.error( "litellm.router.py::async_function_with_fallbacks() - " - "Error occurred while trying to do fallbacks - %s\n%s\n" + "Error occurred while trying to do fallbacks - %s\n" "Debug Information:\nCooldown Deployments=%s", fallback_failure_exception_str, - redact_string(traceback.format_exc()), cooldown_info, ) @@ -7162,7 +7168,7 @@ class Router: k, v, ) in kwargs.items(): # log everything in kwargs except the old previous_models value - prevent nesting - if k not in [_metadata_var, "messages", "original_function"]: + if k != _metadata_var and k not in RETRY_BREADCRUMB_EXCLUDED_KWARGS: previous_model[k] = v elif k == _metadata_var and isinstance(v, dict): previous_model[_metadata_var] = {} diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a69509fc37a..f6ced0bb9d2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1730,6 +1730,7 @@ class ComplexityRouter(CustomLogger): routing_decision=self._build_routing_decision( routed_model=routed_model, cause=cause, + tier=self._tier_for_model(routed_model), escalation_keyword=pin_escalation_keyword, escalated=escalated, conversation_continuing=conversation_continuing, @@ -1797,7 +1798,8 @@ class ComplexityRouter(CustomLogger): if user_message is None: verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model") - if not self.config.plugins and self.config.default_model: + default_model_first: Final = not self.config.plugins and self.config.default_model + if default_model_first: # No plugins configured: preserve the pre-existing default_model-first # priority exactly (changing it would be a silent behavior change for # every non-plugin user, not just a security fix). @@ -1809,12 +1811,14 @@ class ComplexityRouter(CustomLogger): routed_model = await self._pick_model_for_tier( ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs ) + fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, routing_decision=self._build_routing_decision( routed_model=routed_model, cause="default_fallback", + tier=fallback_tier, conversation_continuing=conversation_continuing, ), ) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 1c628ef9ee3..6fad2dd31e9 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -7,6 +7,7 @@ if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject from litellm._logging import verbose_logger +from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError from litellm.types.router import CredentialLiteLLMParams @@ -43,6 +44,22 @@ def resolve_model_group_alias(model_group_alias: object, model: str) -> str | No return target +def truncate_fallback_error_detail(detail: str) -> str: + """ + Bound a fallback failure detail before it is logged or appended to an exception message. + + Each level of the fallback walk records the failure of the level below it, so an + untruncated detail carries every nested failure with it and grows superlinearly with + the number of attempted model groups. One deterministic pre-network failure walked + through a small fallback graph is enough to turn that into hundreds of megabytes of + output on the event-loop thread, which starves the process that produced it. + """ + if len(detail) <= ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: + return detail + dropped: Final = len(detail) - ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + return f"{detail[:ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS]}... [truncated {dropped} characters]" + + def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: """ Hash of the credential params, used for mapping the file id to the right model diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 1c6bb52ccb8..ef48ccc821e 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -1,3 +1,6 @@ +import hashlib +import json +from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Any, Final @@ -19,6 +22,52 @@ else: LitellmRouter = Any +def fallback_attempt_key(fallback_target: object) -> str | None: + """ + Identity of one fallback attempt, so the same attempt is never made twice per request. + + A bare model group name and a `{"model": name}` entry describe the same attempt. An + entry carrying anything else describes a different one and keeps its own identity: a + client-side fallback list overrides request params such as `messages`, and the router + re-targets the group that just failed by attaching `_target_order` or + `_excluded_deployment_ids` to select a different set of deployments inside it. The + payload is hashed rather than kept, so a large `messages` override does not make the + request hold a second copy of itself. + + Returns None for a shape with no usable identity, which is never skipped. + """ + if isinstance(fallback_target, str): + return fallback_target + if not isinstance(fallback_target, dict): + return None + model: Final = fallback_target.get("model") + if tuple(fallback_target) == ("model",) and isinstance(model, str): + return model + serialized: Final = json.dumps(fallback_target, sort_keys=True, default=str) + return hashlib.sha256(serialized.encode()).hexdigest() + + +@dataclass(slots=True) +class AttemptedFallbackTargets: + """ + The fallback attempts a single request has already made. + + One instance is created on the first fallback hop and shared by reference for the rest + of the walk, so an attempt made in one branch is not repeated in a sibling branch. + Without it the walk enumerates paths rather than attempts: a fallback graph containing + a cycle retries one deterministic failure once per path through the cycle, and a + client-side fallback list is re-walked at every level of the recursion. + """ + + keys: frozenset[str] = frozenset() + + def __contains__(self, key: str) -> bool: + return key in self.keys + + def record(self, key: str) -> None: + self.keys = self.keys | frozenset((key,)) + + def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool: """ Handles wildcard routing scenario @@ -106,7 +155,14 @@ async def run_async_fallback( fallback_model_group: List[str] of fallback model groups. example: ["gpt-4", "gpt-3.5-turbo"] original_model_group: The original model group. example: "gpt-3.5-turbo" original_exception: The original exception. - **kwargs: Keyword arguments. + **kwargs: Keyword arguments. `attempted_targets` carries the fallback attempts + already made for this request, created on the first hop and shared by reference + for the rest of the walk. A target already in it is skipped, so neither a + fallback graph that loops back on itself nor a client-side fallback list + re-walked at each level can repeat an attempt that has already failed. Identity + comes from `fallback_attempt_key`, so an entry that overrides request params or + re-targets the failed group with a different deployment selection stays distinct + from a bare name. Returns: The response from the successful fallback model group. @@ -120,10 +176,27 @@ async def run_async_fallback( error_from_fallbacks = original_exception fallback_errors = (get_fallback_error_info(original_exception),) + # 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. + carried_targets: Final = kwargs.get("attempted_targets") + attempted: Final = ( + carried_targets if isinstance(carried_targets, AttemptedFallbackTargets) else AttemptedFallbackTargets() + ) + attempted.record(original_model_group) for mg in fallback_model_group: if mg == original_model_group: continue + attempt_key = fallback_attempt_key(mg) + if attempt_key is not None: + if attempt_key in attempted: + verbose_router_logger.info( + "Skipping fallback to model_group = %s, already attempted for this request", + mask_sensitive_structure(mg), + ) + continue + attempted.record(attempt_key) try: # LOGGING kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) @@ -138,6 +211,7 @@ async def run_async_fallback( fallback_depth = fallback_depth + 1 kwargs["fallback_depth"] = fallback_depth kwargs["max_fallbacks"] = max_fallbacks + kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 6b354a39101..bbb6d758814 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -5,6 +5,7 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing_extensions import Required, TypedDict +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( AktoConfigModel, ) @@ -525,6 +526,15 @@ class BedrockGuardrailConfigModel(BaseModel): description="InvokeGuardrailChecks: block when any sensitiveInformation confidenceScore " ">= this value (scores are in [0,1]). Set to null to make PII detection detect-only.", ) + chunk_budget_chars: int = Field( + default=BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + gt=0, + description="ApplyGuardrail: batch size, in characters, used to re-send content after AWS " + "has rejected a request as too large. Requests AWS accepts are always sent in a single " + "call, so this has no effect until a rejection happens. Defaults to 25,000; a batch AWS " + "still rejects is bisected automatically, so this value only trades round trips against " + "batch size and cannot fail a request on its own.", + ) class LakeraV2GuardrailConfigModel(BaseModel): diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index 05537da67d7..90713b270be 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -2,7 +2,28 @@ Type definitions for WebSearch Interception integration. """ -from typing import TypedDict +from typing import Literal, TypedDict + +from pydantic import BaseModel + + +class AnthropicSearchQuery(BaseModel): + """``input`` of an Anthropic ``server_tool_use`` block for a web search.""" + + query: str + + +class AnthropicServerToolUseBlock(BaseModel): + """ + The ``server_tool_use`` block that must accompany a ``web_search_tool_result``. + + Anthropic requires the pair, with a ``srvtoolu_``-prefixed id shared by both. + """ + + type: Literal["server_tool_use"] = "server_tool_use" + id: str + name: Literal["web_search"] = "web_search" + input: AnthropicSearchQuery class WebSearchInterceptionConfig(TypedDict, total=False): diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 1bf8ba513c2..5665aa3277a 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,4 +1,5 @@ import json +from enum import Enum from typing import TYPE_CHECKING, Any, Final, Literal from typing_extensions import Required, TypedDict, override @@ -1096,3 +1097,17 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): # supported subset and strips the field entirely when nothing remains, so # other edit types (e.g. `clear_thinking_20251015`) never reach Bedrock. context_management: dict + + +class BedrockBatchRecordKind(Enum): + """ + Which OpenAI endpoint shape a line of a Bedrock managed-batch JSONL file + carries. Bedrock batch `modelInput` is always the model's InvokeModel / + Converse body, so every non-embedding shape is normalized to Chat + Completions before being handed to the per-provider transformation. + """ + + CHAT = "chat" + TEXT_COMPLETION = "text_completion" + RESPONSES = "responses" + EMBEDDING = "embedding" diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 6c8fb96a729..6626dea6849 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -2,6 +2,7 @@ Types for auto-router management endpoints """ +from collections.abc import Mapping from typing import Final from pydantic import BaseModel, Field, field_validator @@ -120,6 +121,16 @@ class AutoRouterBenchmarkGroup(AutoRouterBenchmarkTotals): router_name: str = Field(description="The auto-router alias requests were sent to") router_type: str = Field(description="complexity, adaptive or quality") + tier_turns: Mapping[str, int] = Field( + default_factory=dict, + description="Turns per tier, keyed by the tier name the routing decision recorded at " + "request time (never re-derived at read time, since the tier-to-model mapping is " + "mutable config). Tier names are scoped to this group's router_type and are not " + "comparable across types: a complexity router reports 'simple'/'medium'/'complex'/" + "'reasoning', a quality router reports its numeric quality tier, and an adaptive router " + "records no tier at all. Turns no tier served (the classifier fell back to default_model) " + "are absent rather than pooled under a sentinel key, so the values may sum to less than turns", + ) class AutoRouterBenchmarksResponse(BaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d97bdc3532f..8d66b624341 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -28,6 +28,9 @@ class BedrockGuardrailUsage(TypedDict, total=False): sensitiveInformationPolicyUnits: int | None sensitiveInformationPolicyFreeUnits: int | None contextualGroundingPolicyUnits: int | None + contentPolicyImageUnits: int | None + automatedReasoningPolicyUnits: int | None + automatedReasoningPolicies: int | None class BedrockGuardrailOutput(TypedDict, total=False): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py index 3991cee8548..37125c4d583 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py @@ -79,6 +79,15 @@ class ZscalerAIGuardConfigModel(GuardrailConfigModel): json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL}, ) + timeout: float | None = Field( + default=None, + description=( + "Timeout for each Zscaler AI Guard API call, in seconds. Must be positive. " + "Raise it if scans fail under load with 'Connection timed out'. " + "Defaults to 5 seconds." + ), + ) + @model_validator(mode="after") def validate_endpoint_configuration(self) -> "ZscalerAIGuardConfigModel": """ diff --git a/litellm/types/router.py b/litellm/types/router.py index 8b4b547bdcc..4280da08cbb 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -17,7 +17,12 @@ from .completion import CompletionRequest from .embedding import EmbeddingRequest from .llms.openai import OpenAIFileObject from .search import SearchProvider -from .utils import CustomPricingLiteLLMParams, ModelResponse, StandardLoggingRoutingDecision +from .utils import ( + CustomPricingLiteLLMParams, + MirroredPricingParams, + ModelResponse, + StandardLoggingRoutingDecision, +) class ConfigurableClientsideParamsCustomAuth(TypedDict): @@ -122,7 +127,7 @@ class UpdateRouterConfig(BaseModel): model_config = ConfigDict(protected_namespaces=()) -class ModelInfo(BaseModel): +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. updated_at: datetime.datetime | None = None @@ -424,14 +429,7 @@ class DeploymentTypedDict(TypedDict, total=False): model_info: dict -SPECIAL_MODEL_INFO_PARAMS = [ - "input_cost_per_token", - "output_cost_per_token", - "input_cost_per_character", - "output_cost_per_character", - "cache_read_input_token_cost", - "cache_creation_input_token_cost", -] +SPECIAL_MODEL_INFO_PARAMS = tuple(MirroredPricingParams.model_fields) class Deployment(BaseModel): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 329b10e72e7..18cf9461648 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3245,10 +3245,23 @@ class StandardCallbackDynamicParams(TypedDict, total=False): litellm_disabled_callbacks: list[str] | None -class CustomPricingLiteLLMParams(BaseModel): - ## CUSTOM PRICING ## +class MirroredPricingParams(BaseModel): + """Pricing overrides that ``Deployment.__init__`` mirrors from ``litellm_params`` + onto ``model_info``, so both blobs hold the same rate. + + Declared once and inherited by both sides of that mirror, so the two can't drift. + """ + input_cost_per_token: float | None = None output_cost_per_token: float | None = None + input_cost_per_character: float | None = None + output_cost_per_character: float | None = None + cache_read_input_token_cost: float | None = None + cache_creation_input_token_cost: float | None = None + + +class CustomPricingLiteLLMParams(MirroredPricingParams): + ## CUSTOM PRICING ## input_cost_per_second: float | None = None output_cost_per_second: float | None = None output_cost_per_second_1080p: float | None = None @@ -3259,7 +3272,6 @@ class CustomPricingLiteLLMParams(BaseModel): # This allows any model_info parameter to be set in litellm_params input_cost_per_token_flex: float | None = None input_cost_per_token_priority: float | None = None - cache_creation_input_token_cost: float | None = None cache_creation_input_token_cost_above_1hr: float | None = None cache_creation_input_token_cost_above_200k_tokens: float | None = None cache_creation_input_token_cost_above_272k_tokens: float | None = None @@ -3268,7 +3280,6 @@ class CustomPricingLiteLLMParams(BaseModel): cache_creation_input_token_cost_flex: float | None = None cache_creation_input_token_cost_priority: float | None = None cache_creation_input_audio_token_cost: float | None = None - cache_read_input_token_cost: float | None = None cache_read_input_token_cost_flex: float | None = None cache_read_input_token_cost_priority: float | None = None cache_read_input_token_cost_above_200k_tokens: float | None = None @@ -3276,7 +3287,6 @@ class CustomPricingLiteLLMParams(BaseModel): cache_read_input_token_cost_above_272k_tokens_priority: float | None = None cache_read_input_token_cost_above_272k_tokens_flex: float | None = None cache_read_input_audio_token_cost: float | None = None - input_cost_per_character: float | None = None input_cost_per_character_above_128k_tokens: float | None = None input_cost_per_audio_token: float | None = None input_cost_per_token_cache_hit: float | None = None @@ -3298,7 +3308,6 @@ class CustomPricingLiteLLMParams(BaseModel): output_cost_per_token_batches: float | None = None output_cost_per_token_flex: float | None = None output_cost_per_token_priority: float | None = None - output_cost_per_character: float | None = None output_cost_per_audio_token: float | None = None output_cost_per_token_above_128k_tokens: float | None = None output_cost_per_token_above_200k_tokens: float | None = None @@ -3479,6 +3488,7 @@ all_litellm_params = ( "user_continue_message", "fallback_depth", "max_fallbacks", + "attempted_targets", "max_budget", "budget_duration", "use_in_pass_through", diff --git a/litellm/utils.py b/litellm/utils.py index cdf1cc3cf23..911de83b785 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5769,7 +5769,7 @@ def json_schema_type(python_type_name: str): return python_to_json_schema_types.get(python_type_name, "string") -def function_to_dict(input_function) -> dict: # noqa: C901 +def function_to_dict(input_function) -> dict: """Using type hints and numpy-styled docstring, produce a dictionary usable for OpenAI function calling diff --git a/osv-scanner.toml b/osv-scanner.toml index 7ab450945f5..4ef612e3a70 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -1,3 +1,13 @@ +[[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/pyproject.toml b/pyproject.toml index 414b09eb3b4..35fd949c2e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,8 +66,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.83", - "litellm-enterprise==0.1.53", + "litellm-proxy-extras==0.4.84", + "litellm-enterprise==0.1.54", "RestrictedPython>=8.1,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 60356eda05b..fdc81fac196 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3121 + "limit": 3114 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 834 }, "ANN201": { - "limit": 2033 + "limit": 2031 }, "ANN202": { "limit": 865 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1630 + "limit": 1555 }, "ASYNC230": { "limit": 11 @@ -56,11 +56,8 @@ "B026": { "limit": 3 }, - "B033": { - "limit": 0 - }, "BLE001": { - "limit": 2926 + "limit": 2924 }, "C401": { "limit": 8 @@ -81,7 +78,7 @@ "limit": 1 }, "C901": { - "limit": 315 + "limit": 314 }, "D419": { "limit": 6 @@ -113,18 +110,6 @@ "F401": { "limit": 17 }, - "FURB136": { - "limit": 0 - }, - "FURB168": { - "limit": 0 - }, - "FURB188": { - "limit": 0 - }, - "I001": { - "limit": 0 - }, "LOG015": { "limit": 5 }, @@ -135,20 +120,11 @@ "limit": 27 }, "PERF401": { - "limit": 23 - }, - "PERF402": { - "limit": 0 + "limit": 12 }, "PERF403": { "limit": 34 }, - "PIE790": { - "limit": 0 - }, - "PIE800": { - "limit": 0 - }, "PIE804": { "limit": 18 }, @@ -158,9 +134,6 @@ "PLC0206": { "limit": 26 }, - "PLC0208": { - "limit": 0 - }, "PLC0414": { "limit": 46 }, @@ -170,24 +143,12 @@ "PLR0206": { "limit": 1 }, - "PLR0402": { - "limit": 0 - }, "PLR1704": { "limit": 3 }, - "PLR1711": { - "limit": 0 - }, "PLR1714": { "limit": 257 }, - "PLR1730": { - "limit": 0 - }, - "PLR2044": { - "limit": 0 - }, "PLW0127": { "limit": 57 }, @@ -206,27 +167,12 @@ "PLW1510": { "limit": 2 }, - "PYI030": { - "limit": 0 - }, "PYI036": { "limit": 3 }, - "PYI041": { - "limit": 0 - }, - "PYI064": { - "limit": 0 - }, - "RET501": { - "limit": 0 - }, "RET504": { "limit": 177 }, - "RUF010": { - "limit": 0 - }, "RUF012": { "limit": 241 }, @@ -236,23 +182,14 @@ "RUF019": { "limit": 38 }, - "RUF022": { - "limit": 0 - }, - "RUF023": { - "limit": 0 - }, "RUF046": { "limit": 4 }, - "RUF051": { - "limit": 0 - }, "RUF059": { "limit": 67 }, "RUF100": { - "limit": 100 + "limit": 0 }, "S110": { "limit": 218 @@ -272,18 +209,12 @@ "SIM113": { "limit": 3 }, - "SIM114": { - "limit": 0 - }, "SIM115": { "limit": 2 }, "SIM117": { "limit": 7 }, - "SIM118": { - "limit": 0 - }, "SIM201": { "limit": 1 }, @@ -302,11 +233,8 @@ "TC004": { "limit": 5 }, - "TC005": { - "limit": 0 - }, "TID251": { - "limit": 1240 + "limit": 1238 }, "TRY002": { "limit": 528 @@ -323,46 +251,13 @@ "TRY300": { "limit": 860 }, - "UP006": { - "limit": 0 - }, - "UP007": { - "limit": 0 - }, - "UP008": { - "limit": 0 - }, - "UP012": { - "limit": 0 - }, - "UP018": { - "limit": 0 - }, - "UP024": { - "limit": 0 - }, "UP028": { "limit": 2 }, "UP031": { "limit": 2 }, - "UP032": { - "limit": 0 - }, - "UP034": { - "limit": 0 - }, - "UP035": { - "limit": 0 - }, "UP036": { "limit": 1 - }, - "UP037": { - "limit": 0 - }, - "UP045": { - "limit": 0 } } diff --git a/ruff-strict.toml b/ruff-strict.toml index 01faf04805f..974c49c787b 100644 --- a/ruff-strict.toml +++ b/ruff-strict.toml @@ -4,6 +4,17 @@ extend = "ruff.toml" preview = true select = ["ANN", "ASYNC230", "B004", "B006", "B008", "B009", "B010", "B018", "B019", "B021", "B026", "B033", "BLE", "C401", "C404", "C405", "C408", "C414", "C419", "C901", "D419", "DTZ001", "DTZ003", "DTZ005", "DTZ006", "DTZ007", "DTZ011", "EXE001", "EXE002", "F401", "FURB136", "FURB168", "FURB188", "I001", "LOG015", "N999", "PERF102", "PERF401", "PERF402", "PERF403", "PIE790", "PIE800", "PIE804", "PIE810", "PLC0206", "PLC0208", "PLC0414", "PLR0124", "PLR0206", "PLR0402", "PLR1704", "PLR1711", "PLR1714", "PLR1730", "PLR2044", "PLW0127", "PLW0133", "PLW0602", "PLW0603", "PLW1508", "PLW1510", "PYI030", "PYI036", "PYI041", "PYI064", "RET501", "RET504", "RUF010", "RUF012", "RUF015", "RUF019", "RUF022", "RUF023", "RUF046", "RUF051", "RUF059", "RUF100", "S110", "S112", "SIM101", "SIM102", "SIM103", "SIM113", "SIM114", "SIM115", "SIM117", "SIM118", "SIM201", "SIM210", "SIM211", "SIM222", "SIM401", "TC004", "TC005", "TID251", "TRY002", "TRY004", "TRY201", "TRY203", "TRY300", "UP006", "UP007", "UP008", "UP012", "UP018", "UP024", "UP028", "UP031", "UP032", "UP034", "UP035", "UP036", "UP037", "UP045"] extend-select = [] +# Overrides the inherited list: rules this gate enforces itself must NOT be external here, +# so this config's RUF100 flags their stale `# noqa` directives. What remains external is +# only what other tooling enforces: every base ruff.toml rule this select list doesn't +# re-enable (all of the default E/F families plus T20/PGH004/RUF008/RUF009, minus the +# strict-selected F401 and RUF100; F4 is split out so stale F401 noqas stay detectable), +# plus upstream litellm's ruff config. +external = [ + "T20", "PGH004", "RUF008", "RUF009", "E4", "E7", "E9", + "F402", "F404", "F406", "F407", "F5", "F6", "F7", "F8", "F9", + "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", +] [lint.mccabe] max-complexity = 15 diff --git a/ruff.toml b/ruff.toml index 095e3e24c52..9b90910b355 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,12 +1,26 @@ lint.ignore = ["F405", "E402", "F403"] -lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] +# The second group is the strict gate's graduates: rules the codebase already has zero +# violations of, so they hard-fail here instead of being ratcheted in ruff-strict-budget.json. +# That gives editors and `ruff check --fix` the diagnostic, which the gate script cannot. +lint.extend-select = [ + "T20", "PGH004", "RUF008", "RUF009", "RUF100", + "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208", + "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PYI030", "PYI041", "PYI064", "RET501", "RUF010", + "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", "UP012", + "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", +] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip # `# noqa` directives that protect rules enforced elsewhere. List those codes as external # so RUF100 leaves their directives alone: the strict gate (ruff-strict.toml) and upstream # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ - # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "C901", "TID251", + # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml). + # Family entries whose every strict rule graduated into extend-select above (FURB), and + # standalone graduated codes (I001, RUF010, RUF022, RUF023, RUF051), are dropped so this + # config's RUF100 polices their directives itself. + "ANN", "ASYNC230", "B", "C", "D419", "DTZ", "EXE", "LOG015", "N999", "PERF", + "PIE", "PL", "PYI", "RET", "RUF012", "RUF015", "RUF019", + "RUF046", "RUF059", "S110", "S112", "SIM", "TC", "TID251", "TRY", "UP", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/schema.prisma b/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1439,6 +1439,7 @@ model LiteLLM_AutoRouterSession { total_tokens BigInt @default(0) spend Float @default(0) saved_spend Float @default(0) + tier_turns Json @default("{}") @@id([api_key, session_id, router_name]) @@index([last_turn_at], map: "idx_autorouter_session_last_turn") diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 10a78483643..e97cd1bca00 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -10,7 +10,10 @@ content at the merge-base with the target branch and fails (exits 1, red) if: * a rule was dropped from a budget (its ceiling effectively became infinite), or * an entire budget file was deleted. -New rules and lowered/equal limits are fine. +New rules and lowered/equal limits are fine. So is a rule that graduated: once a +paired config (ruff.toml for the ruff-strict budget) selects the rule outright it +hard-fails at the first violation, which is stricter than any ceiling the budget +could hold, so dropping its entry tightens the guard rather than removing it. This is deliberately NOT a gating check. It should turn the run red so that a loosening is impossible to miss in review, but it must stay OUT of the @@ -30,7 +33,9 @@ import argparse import json import subprocess import sys +import tomllib from pathlib import Path +from types import MappingProxyType from typing import NamedTuple REPO_ROOT = Path(__file__).resolve().parent.parent @@ -40,6 +45,7 @@ DEFAULT_BUDGETS: tuple[str, ...] = ( "type-discipline-budget.json", "basedpyright-code-budget.json", ) +GRADUATION_CONFIGS = MappingProxyType({"ruff-strict-budget.json": "ruff.toml"}) class Regression(NamedTuple): @@ -106,24 +112,57 @@ def _limits(budget: dict) -> dict[str, int]: } +def selectors_hard_failed_by(lint: dict) -> tuple[str, ...]: + """A ruff `[lint]` table's selected codes, minus anything `ignore` turns back off. + + `lint.ignore` wins over `lint.extend-select` in ruff, so an ignored code is not + actually enforced and must not count as a graduation. + """ + ignored = tuple(lint.get("ignore", ())) + return tuple( + selector + for selector in lint.get("extend-select", ()) + if not (ignored and selector.startswith(ignored)) + ) + + +def graduated_selectors(rel: str) -> tuple[str, ...]: + """Selectors the budget's paired ruff config hard-fails, so its ceiling is moot.""" + config = GRADUATION_CONFIGS.get(rel) + if config is None or not (REPO_ROOT / config).exists(): + return () + return selectors_hard_failed_by( + tomllib.loads((REPO_ROOT / config).read_text()).get("lint", {}) + ) + + def _regression_detail( rule: str, base_limits: dict[str, int], head_limits: dict[str, int], + graduated: tuple[str, ...], ) -> str | None: - """Why `rule` regressed vs base, or None when it held flat or fell. + """Why `rule` regressed vs base, or None when it held flat, fell, or graduated. - A dropped rule is terminal; otherwise the only loosening left is a raised limit. + A dropped rule is terminal unless it graduated; otherwise the only loosening + left is a raised limit. """ base_limit = base_limits[rule] if rule not in head_limits: + if graduated and rule.startswith(graduated): + return None return f"rule dropped (limit {base_limit} -> removed)" if head_limits[rule] > base_limit: return f"limit raised {base_limit} -> {head_limits[rule]}" return None -def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regression]: +def regressions_for( + rel: str, + base: dict | None, + head: dict | None, + graduated: tuple[str, ...] = (), +) -> list[Regression]: if base is None: return [] # new budget file: nothing to ratchet against yet if head is None: @@ -133,7 +172,7 @@ def regressions_for(rel: str, base: dict | None, head: dict | None) -> list[Regr return [ Regression(rel, rule, detail) for rule in sorted(base_limits) - if (detail := _regression_detail(rule, base_limits, head_limits)) is not None + if (detail := _regression_detail(rule, base_limits, head_limits, graduated)) is not None ] @@ -164,7 +203,7 @@ def main() -> int: print(f"skip {rel}: new file (no base at {args.base} to ratchet against)") continue checked.append(rel) - regressions.extend(regressions_for(rel, base, head)) + regressions.extend(regressions_for(rel, base, head, graduated_selectors(rel))) if regressions: print( diff --git a/scripts/install_git_hooks.sh b/scripts/install_git_hooks.sh index 7ea8c3ff2e9..8f4e79e4ddd 100755 --- a/scripts/install_git_hooks.sh +++ b/scripts/install_git_hooks.sh @@ -35,7 +35,7 @@ These hooks enforce Conventional Commits and Conventional Branches. Bypass with --no-verify when you need to (e.g. for emergency hotfixes). The CI-equivalent lint is deliberately not installed as an auto-firing hook -(it can take minutes); run it on demand with 'make pre-commit' before committing. +(it can take minutes); run it on demand with 'make check' before committing. To uninstall: git config --unset core.hooksPath EOF diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh index af8335e0e84..afe55603466 100755 --- a/scripts/pre_commit_lint.sh +++ b/scripts/pre_commit_lint.sh @@ -1,18 +1,25 @@ #!/usr/bin/env bash # -# pre_commit_lint.sh — shift CI lint left. Run it (via `make pre-commit`) right -# before `git commit`; it inspects your staged files and runs only the matching -# gating CI checks, so a clean run means a green CI lint: -# - litellm/ Python staged -> `make lint` (test-linting.yml's lint job) -# - tests/e2e Python staged -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) -# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) -# - dashboard staged -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) -# - proxy/types staged -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml) +# pre_commit_lint.sh — shift CI lint left. Run it (via `make check`, formerly +# `make pre-commit`) before `git commit`, or after committing (e.g. a merge +# commit) to predict CI for the branch. It picks the files in scope and runs +# only the matching gating CI checks, so a clean run means a green CI lint: +# - anything staged -> scope is the staged files; changed-but-unstaged files +# whose checks were skipped are called out +# - nothing staged -> scope is the working tree's diff against the merge base +# with origin/litellm_internal_staging, untracked files included +# The per-area checks: +# - litellm/ Python -> `make lint` (test-linting.yml's lint job) +# - tests/e2e Python -> `make lint-e2e-basedpyright` (test-linting.yml's e2e type-check step) +# + raw HTTP client ban (test-code-quality.yml's check_e2e_no_raw_requests) +# - dashboard -> prettier + eslint + lint budgets (test-litellm-ui-build.yml's frontend-lint) +# - proxy/types -> regenerate dashboard API types and fail on drift (check-ui-api-types.yml) # -# Each block is skipped when no matching files are staged, so unrelated commits stay -# fast. This is intentionally not auto-installed as a git hook (see scripts/install_git_hooks.sh): -# the dashboard and basedpyright passes can take minutes, so it's run on demand rather -# than firing on every human commit. It is hook-compatible if you want that anyway: +# Each block is skipped when no matching files are in scope, so unrelated commits +# stay fast. This is intentionally not auto-installed as a git hook (see +# scripts/install_git_hooks.sh): the dashboard and basedpyright passes can take +# minutes, so it's run on demand rather than firing on every human commit. It is +# hook-compatible if you want that anyway: # `ln -s ../../scripts/pre_commit_lint.sh .git/hooks/pre-commit`. set -eu @@ -20,56 +27,105 @@ set -eu if [ -z "${PRE_COMMIT_LINT_INNER:-}" ]; then log_file=$(git rev-parse --path-format=absolute --git-path pre_commit_lint.log) if : > "$log_file" 2>/dev/null; then - echo "pre-commit: logging full output to $log_file" + echo "check: logging full output to $log_file" PRE_COMMIT_LINT_INNER=1 "$0" "$@" 2>&1 | tee "$log_file" pipe_status=("${PIPESTATUS[@]}") if [ "${pipe_status[1]}" -eq 0 ]; then - echo "pre-commit: full log: $log_file" + echo "check: full log: $log_file" else - echo "pre-commit: WARNING - writing $log_file failed; the log may be incomplete" >&2 + echo "check: WARNING - writing $log_file failed; the log may be incomplete" >&2 fi exit "${pipe_status[0]}" fi - echo "pre-commit: WARNING - cannot write $log_file; output will not be saved" >&2 + echo "check: WARNING - cannot write $log_file; output will not be saved" >&2 PRE_COMMIT_LINT_INNER=1 exec "$0" "$@" fi repo_root=$(git rev-parse --show-toplevel) cd "$repo_root" -staged=$(git diff --cached --name-only --diff-filter=ACMR) -staged_match() { printf '%s\n' "$staged" | grep -E "$1" || true; } +staged=$(git diff --cached --name-only --diff-filter=ACMRD) +unstaged=$(git diff --name-only) +untracked=$(git ls-files --others --exclude-standard) + +if [ -n "$staged" ]; then + scope=$staged +else + git fetch --quiet origin litellm_internal_staging 2>/dev/null || true + merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || { + echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2 + echo " Fix: git fetch origin litellm_internal_staging" >&2 + exit 1 + } + scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u) + if [ -z "$scope" ]; then + echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)" + exit 0 + fi + echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:" + printf '%s\n' "$scope" | sed 's/^/ /' +fi + +scope_match() { printf '%s\n' "$scope" | grep -E "$1" || true; } + +existing_files() { + while IFS= read -r f; do + if [ -f "$f" ]; then printf '%s\n' "$f"; fi + done +} + +litellm_py_pattern='^litellm/.*\.py$' +e2e_py_pattern='^tests/e2e/.*\.py$' +spec_pattern='^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$' +ui_prettier_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$' +ui_eslint_pattern='^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$' # CI's lint job (test-linting.yml) only inspects litellm/, so a tests-only or # scripts-only commit can't turn it red; scope the trigger there to skip the slow # make lint when it couldn't catch anything. -litellm_py_files=$(staged_match '^litellm/.*\.py$') -e2e_py_files=$(staged_match '^tests/e2e/.*\.py$') +litellm_py_files=$(scope_match "$litellm_py_pattern") +e2e_py_files=$(scope_match "$e2e_py_pattern") # ruff format (and CI's format step) skip enterprise; the rest of make lint covers it. -fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' || true) +fmt_files=$(printf '%s\n' "$litellm_py_files" | grep -v '^litellm/enterprise/' | existing_files) # check-ui-api-types.yml triggers on any file under litellm/proxy or litellm/types # (Prisma schema and configs included, not just Python) plus the generator and its # lockfiles, so match that whole trigger set rather than a Python subset. -spec_files=$(staged_match '^(litellm/(proxy|types)/.*|ui/litellm-dashboard/(scripts/gen-api-types\.mjs|package\.json|package-lock\.json|src/lib/http/schema\.d\.ts))$') +spec_files=$(scope_match "$spec_pattern") # CI's frontend-lint runs prettier over a wider extension set than eslint; keep that # split so this flags exactly what the job would. -ui_prettier_files=$(staged_match '^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs|json|css|scss|md|mdx|yml|yaml|html)$') -ui_eslint_files=$(staged_match '^ui/litellm-dashboard/.*\.(js|jsx|ts|tsx|mjs|cjs)$') +ui_prettier_changed=$(scope_match "$ui_prettier_pattern") +ui_eslint_changed=$(scope_match "$ui_eslint_pattern") +ui_prettier_files=$(printf '%s\n' "$ui_prettier_changed" | existing_files) +ui_eslint_files=$(printf '%s\n' "$ui_eslint_changed" | existing_files) -# CI lints the committed tree, so this script predicts CI for what you have STAGED -# (every trigger above reads `git diff --cached`). The tools it runs, though, read -# the working tree, so unstaged edits to tracked files and untracked files fold -# into the result and a green/red here won't match a commit of just the staged -# changes. There's no safe way to lint the index in place, so surface the gap -# instead of hiding it: stage everything you intend to commit before trusting a -# pass. This only warns; it never blocks or touches your changes. -unstaged=$(git diff --name-only) -untracked=$(git ls-files --others --exclude-standard) -if [ -n "$unstaged" ] || [ -n "$untracked" ]; then - echo "pre-commit: NOTE - unstaged/untracked changes are included in these checks but" >&2 - echo " won't be in a commit of only your staged changes, so this result may differ from" >&2 - echo " CI. Stage everything you intend to commit (git add) for an accurate prediction:" >&2 - printf '%s\n' "$unstaged" "$untracked" | sed '/^$/d' | sed 's/^/ /' >&2 +# CI lints the committed tree, so with staged files this script predicts CI for +# what you have STAGED (every trigger above reads `git diff --cached`). The tools +# it runs, though, read the working tree, so unstaged edits to tracked files and +# untracked files fold into the result and a green/red here won't match a commit +# of just the staged changes. There's no safe way to lint the index in place, so +# surface the gap instead of hiding it: stage everything you intend to commit +# before trusting a pass. This only warns; it never blocks or touches your changes. +if [ -n "$staged" ]; then + not_staged=$(printf '%s\n' "$unstaged" "$untracked" | sed '/^$/d' | sort -u) + if [ -n "$not_staged" ]; then + echo "check: NOTE - unstaged/untracked changes are included in these checks but" >&2 + echo " won't be in a commit of only your staged changes, so this result may differ from" >&2 + echo " CI. Stage everything you intend to commit (git add) for an accurate prediction:" >&2 + printf '%s\n' "$not_staged" | sed 's/^/ /' >&2 + fi + warn_skipped() { + local check_name=$1 pattern=$2 triggered=$3 + [ -n "$triggered" ] && return 0 + local missed + missed=$(printf '%s\n' "$not_staged" | grep -E "$pattern" || true) + [ -z "$missed" ] && return 0 + echo "check: SKIPPED $check_name because these changed files are not staged:" >&2 + printf '%s\n' "$missed" | sed 's/^/ /' >&2 + } + warn_skipped "Python lint (make lint)" "$litellm_py_pattern" "$litellm_py_files" + warn_skipped "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_pattern" "$e2e_py_files" + warn_skipped "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_pattern" "$ui_prettier_changed" + warn_skipped "dashboard API-type sync (npm run gen:api)" "$spec_pattern" "$spec_files" fi lint_dashboard() { @@ -114,15 +170,15 @@ bootstrap_hint() { python_checks() { local rc=0 - echo "pre-commit: linting Python (make lint)" - make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make pre-commit." >&2; rc=1; } + echo "check: linting Python (make lint)" + make lint || { echo "✗ Python lint failed. Fix the reds above, then re-run make check." >&2; rc=1; } # `make lint` format-checks files in origin/base...HEAD, which at pre-commit time - # predates the staged change, so format-check the staged litellm files directly to + # predates the staged change, so format-check the scoped litellm files directly to # cover a brand-new commit before it lands. if [ -n "$fmt_files" ]; then - echo "pre-commit: ruff format --check (staged litellm files)" + echo "check: ruff format --check (scoped litellm files)" printf '%s\n' "$fmt_files" | xargs uv run --no-sync ruff format --check --exclude '/enterprise/' \ - || { echo "✗ Unformatted staged files. Fix with: make format, then re-stage." >&2; rc=1; } + || { echo "✗ Unformatted files in scope. Fix with: make format, then re-stage." >&2; rc=1; } fi return $rc } @@ -146,18 +202,18 @@ if [ -n "$litellm_py_files" ]; then fi if [ -n "$e2e_py_files" ] && [ -z "$litellm_py_files" ]; then - echo "pre-commit: type-checking tests/e2e (make lint-e2e-basedpyright)" - make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make pre-commit." >&2; status=1; } + echo "check: type-checking tests/e2e (make lint-e2e-basedpyright)" + make lint-e2e-basedpyright || { echo "✗ tests/e2e basedpyright failed. Fix the errors above, then re-run make check." >&2; status=1; } fi if [ -n "$e2e_py_files" ]; then - echo "pre-commit: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)" + echo "check: checking tests/e2e raw HTTP client ban (check_e2e_no_raw_requests)" uv run --no-sync python tests/code_coverage_tests/check_e2e_no_raw_requests.py \ - || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make pre-commit." >&2; status=1; } + || { echo "✗ Raw HTTP client import in tests/e2e. Route the call through tests/e2e/e2e_http.py, then re-run make check." >&2; status=1; } fi dashboard_checks() { - echo "pre-commit: linting dashboard (prettier + eslint + lint budgets)" + echo "check: linting dashboard (prettier + eslint + lint budgets)" if [ ! -d ui/litellm-dashboard/node_modules ]; then echo "✗ ui/litellm-dashboard/node_modules is missing; dashboard lint cannot run." >&2 bootstrap_hint @@ -166,7 +222,7 @@ dashboard_checks() { lint_dashboard || { echo "✗ Dashboard lint failed. See above; format with: (cd ui/litellm-dashboard && npm run format)." >&2; return 1; } } -if [ -n "$ui_prettier_files" ] || [ -n "$ui_eslint_files" ]; then +if [ -n "$ui_prettier_changed" ] || [ -n "$ui_eslint_changed" ]; then dash_log=$(mktemp) set -m dashboard_checks > "$dash_log" 2>&1 & @@ -176,7 +232,7 @@ fi genapi_checks() { local status=0 - echo "pre-commit: checking dashboard API types are in sync (npm run gen:api)" + echo "check: checking dashboard API types are in sync (npm run gen:api)" # gen-api-types.mjs imports litellm.proxy.proxy_server, which needs the proxy deps # and an up-to-date Prisma client; check-ui-api-types.yml installs those and runs # prisma generate before gen:api, so mirror that here or a stale client can mask @@ -194,7 +250,7 @@ genapi_checks() { status=1 elif ( cd ui/litellm-dashboard && LITELLM_PYTHON="uv run --no-sync python" npm run gen:api ); then if ! git diff --quiet -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then - echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make pre-commit only if other checks failed too." >&2 + echo "✗ Dashboard API types are stale; regenerated src/lib/http/schema.d.ts. Stage it and commit; re-run make check only if other checks failed too." >&2 status=1 fi else diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 101519c0b08..7c744f04064 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -7,13 +7,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [0.4.0] - 2026-08-06 + ### Fixed - **organization**: Send `PATCH` instead of `POST` to `/organization/update` and `/organization/member_update`, matching the methods the LiteLLM proxy serves; organization and organization member updates previously failed with a 405 +- **team_member**: Include `role` in the update payload so a role change on an existing `litellm_team_member` is applied instead of being silently dropped ### Changed - The provider source of truth moved to `terraform/provider/` in [BerriAI/litellm](https://github.com/BerriAI/litellm); this repository is now a release mirror. CI in the monorepo statically audits every endpoint the provider calls against the proxy's OpenAPI schema on every change +- **mcp_server**, **vector_store**: `env` and `litellm_params` are now marked sensitive, so they are redacted from plan/apply output, and they are no longer read back from the API into state — the configured value is authoritative. If the proxy returns values that differ from the configuration, that drift is no longer surfaced on refresh +- Dependency updates: `grpc` and `golang.org/x` modules + +## [0.3.0] - 2026-07-13 + +Released from the mirror repository before the source move was complete; this entry backfills it in the monorepo changelog. + +### Added + +- **model**: Add optional `pricing_base_model` attribute that sets `model_info.base_model` (the cost-map lookup key) independently of routing. Deployments whose routing name differs from the pricing key (for example Azure Data Zone, routed as `azure/gpt-4.1` but priced via `us/gpt-4.1-2025-04-14`) can now be billed correctly without breaking routing. When unset, behavior is unchanged and `base_model` continues to drive both routing and pricing (#47) ## [0.2.2] - 2026-05-13 diff --git a/terraform/provider/docs/resources/model.md b/terraform/provider/docs/resources/model.md index 5a46fe2f073..0409b48b391 100644 --- a/terraform/provider/docs/resources/model.md +++ b/terraform/provider/docs/resources/model.md @@ -118,6 +118,8 @@ The following arguments are supported: * `base_model` - (Required) string. The actual model identifier from the provider (e.g., "gpt-4", "claude-2"). +* `pricing_base_model` - (Optional) string. A pricing key fed to `model_info.base_model` **independently of routing**. When set, `litellm_params.model` still routes via `base_model`, but LiteLLM looks up cost against this key. Useful when the routing/deployment name differs from the cost-map key — e.g. an Azure deployment routed as `azure/gpt-4.1` whose real tier is Data Zone: set `pricing_base_model = "us/gpt-4.1-2025-04-14"` so it is billed at the Data Zone rate. When unset, `base_model` drives pricing as before. + * `litellm_credential_name` - (Optional) string. Name of a LiteLLM credential to use for this model. * `tier` - (Optional) string. The usage tier for this model. Valid values are `"free"` or `"paid"`. Default: `"free"`. diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index 2858b6e763d..4bad057871d 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -73,6 +73,14 @@ func resourceLiteLLMModel() *schema.Resource { Type: schema.TypeString, Required: true, }, + "pricing_base_model": { + // Optional pricing key fed to model_info.base_model, DECOUPLED + // from routing. When set, litellm_params.model still routes via + // base_model, but cost is looked up against this key (e.g. + // "us/gpt-4.1-2025-04-14" for Azure Data Zone pricing). + Type: schema.TypeString, + Optional: true, + }, "tier": { Type: schema.TypeString, Optional: true, diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 40766c8e312..fc5d5b09dd5 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -68,6 +68,14 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e baseModel := d.Get("base_model").(string) modelName := fmt.Sprintf("%s/%s", customLLMProvider, baseModel) + // Pricing base_model, decoupled from routing. When pricing_base_model is + // set it feeds model_info.base_model (the cost-lookup key) WITHOUT changing + // the routing string above; otherwise base_model drives pricing as before. + pricingBaseModel := baseModel + if v, ok := d.GetOk("pricing_base_model"); ok && v.(string) != "" { + pricingBaseModel = v.(string) + } + // Generate a UUID for new models modelID := d.Id() if !isUpdate { @@ -240,7 +248,7 @@ func createOrUpdateModel(d *schema.ResourceData, m interface{}, isUpdate bool) e ModelInfo: ModelInfo{ ID: modelID, DBModel: true, - BaseModel: baseModel, + BaseModel: pricingBaseModel, Tier: d.Get("tier").(string), Mode: d.Get("mode").(string), TeamID: d.Get("team_id").(string), @@ -306,7 +314,16 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { d.Set("rpm", GetIntValue(modelResp.LiteLLMParams.RPM, d.Get("rpm").(int))) d.Set("model_api_base", GetStringValue(modelResp.LiteLLMParams.APIBase, d.Get("model_api_base").(string))) d.Set("api_version", GetStringValue(modelResp.LiteLLMParams.APIVersion, d.Get("api_version").(string))) - d.Set("base_model", GetStringValue(modelResp.ModelInfo.BaseModel, d.Get("base_model").(string))) + // base_model / pricing_base_model read-back. When pricing_base_model is + // configured, model_info.base_model holds the PRICING key, so recover the + // routing base_model from state (not returned by the API) and read + // pricing_base_model from model_info. + if pbm, ok := d.GetOk("pricing_base_model"); ok && pbm.(string) != "" { + d.Set("base_model", d.Get("base_model").(string)) + d.Set("pricing_base_model", GetStringValue(modelResp.ModelInfo.BaseModel, pbm.(string))) + } else { + d.Set("base_model", GetStringValue(modelResp.ModelInfo.BaseModel, d.Get("base_model").(string))) + } d.Set("tier", GetStringValue(modelResp.ModelInfo.Tier, d.Get("tier").(string))) d.Set("mode", GetStringValue(modelResp.ModelInfo.Mode, d.Get("mode").(string))) d.Set("team_id", GetStringValue(modelResp.ModelInfo.TeamID, d.Get("team_id").(string))) diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 5999e5772d1..dc69bd42171 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -138,7 +138,7 @@ Mark live tests with `@pytest.mark.e2e` (on the class or the module). Pure cover Before you push -1. Run `make lint-e2e-basedpyright` (or `make pre-commit` with your changes staged); the harness is fully typed and the gate allows zero basedpyright errors, enforced in CI on any PR touching `tests/e2e/**/*.py` +1. Run `make lint-e2e-basedpyright` (or `make check` with your changes staged); the harness is fully typed and the gate allows zero basedpyright errors, enforced in CI on any PR touching `tests/e2e/**/*.py` 2. Add the models your test needs to the config your local proxy loads diff --git a/tests/e2e/a2a/a2a_client.py b/tests/e2e/a2a/a2a_client.py index e83897025a3..605dd8fb7e5 100644 --- a/tests/e2e/a2a/a2a_client.py +++ b/tests/e2e/a2a/a2a_client.py @@ -17,6 +17,7 @@ from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field +from e2e_config import settle_propagation from e2e_http import NoBody, Result, Success, get_external, is_ok from proxy_client import ProxyClient @@ -298,7 +299,9 @@ class A2AClient: the next DB reload. A card read or message/send issued the instant this returns can therefore 404 on the agent it just created. Waiting here keeps every caller from having to poll, the same way ProxyClient.create_model - waits for a new model to become servable. + waits for a new model to become servable -- including the settle that + covers the other replicas, since one successful card read only proves the + replica that answered it has the agent. """ result = self.proxy.transport.post( "/v1/agents", @@ -307,7 +310,9 @@ class A2AClient: response_type=AgentResponse, ) if isinstance(result, Success): + written_at = time.monotonic() self._await_agent_servable(result.data.agent_id) + settle_propagation(written_at) return result def _await_agent_servable(self, agent_id: str) -> None: diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index 34770bc596b..277478eebaf 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -7,6 +7,7 @@ environment so the same tests run against localhost or a deployed proxy. from __future__ import annotations import os +import time import uuid from pathlib import Path @@ -75,6 +76,18 @@ POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120")) POLL_INTERVAL = float(os.environ.get("E2E_POLL_INTERVAL", "5")) REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60")) +# How long a control-plane write (/model/new, /guardrails, /v1/agents) may take to +# reach EVERY replica. Distinct from POLL_TIMEOUT, which is sized for spend-row +# flush; this one is sized for the proxy's config reload +# (`proxy_config_reload_interval_seconds`, 30s by default and 7s on the e2e stack) +# plus margin. +# +# The barriers below wait this out instead of returning on first sight, because a +# single successful read only proves ONE replica converged: every request opens a +# fresh connection, so a load-balanced Service routes each one independently and +# the next call re-rolls. See ProxyClient._await_model_servable. +PROPAGATION_TIMEOUT = float(os.environ.get("E2E_PROPAGATION_TIMEOUT", "15")) + EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes") # Deliberately modest concurrency. The suite shares its proxy with every other @@ -137,3 +150,16 @@ def unique_marker() -> str: """A short unique token per call/run, so concurrent runs and the shared response cache never collide on prompts, tags, or customer ids.""" return uuid.uuid4().hex[:12] + + +def settle_propagation(written_at: float) -> None: + """Block until PROPAGATION_TIMEOUT has elapsed since `written_at`, a + `time.monotonic()` stamp taken the moment a control-plane write returned. + + Callers that already polled for the object still need this: the poll proves one + replica has it, not all of them. Waiting out the config-reload budget is what + makes the object safe to use on whichever replica the next request lands on. + """ + remaining = PROPAGATION_TIMEOUT - (time.monotonic() - written_at) + if remaining > 0: + time.sleep(remaining) diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index 93861d19922..85964529ada 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -11,7 +11,7 @@ from typing import Literal from pydantic import BaseModel -from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, unique_marker +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation, unique_marker from e2e_http import NoBody, Result, Success, unwrap from lifecycle import ResourceManager from models import ( @@ -104,25 +104,14 @@ class GuardrailsClient: proxy: ProxyClient def create_content_filter_guardrail(self, name: str, blocked_keyword: str) -> str: - return unwrap( - self.proxy.transport.post( - "/guardrails", - headers=self.proxy.transport.master, - json=GuardrailCreateBody( - guardrail=GuardrailSpecBody( - guardrail_name=name, - litellm_params=ContentFilterParamsBody( - mode="pre_call", - default_on=True, - blocked_words=[ - BlockedWordBody(keyword=blocked_keyword, action="BLOCK") - ], - ), - ) - ), - response_type=GuardrailCreateResponse, - ) - ).guardrail_id + return self.register( + name, + ContentFilterParamsBody( + mode="pre_call", + default_on=True, + blocked_words=[BlockedWordBody(keyword=blocked_keyword, action="BLOCK")], + ), + ) def create_bedrock_guardrail( self, @@ -141,24 +130,15 @@ class GuardrailsClient: test takes out whatever else is running. Callers select the guardrail per-request instead, which keeps the blast radius to the test that wants it. """ - return unwrap( - self.proxy.transport.post( - "/guardrails", - headers=self.proxy.transport.master, - json=GuardrailCreateBody( - guardrail=GuardrailSpecBody( - guardrail_name=name, - litellm_params=BedrockGuardrailParamsBody( - mode="pre_call", - default_on=default_on, - guardrailIdentifier=identifier, - guardrailVersion=version, - ), - ) - ), - response_type=GuardrailCreateResponse, - ) - ).guardrail_id + return self.register( + name, + BedrockGuardrailParamsBody( + mode="pre_call", + default_on=default_on, + guardrailIdentifier=identifier, + guardrailVersion=version, + ), + ) def create_backend_model(self, resources: ResourceManager, prefix: str = "e2e-guard-backend") -> str: """Register a gemini chat deployment for a guardrail test to run against @@ -174,11 +154,19 @@ class GuardrailsClient: return model_name def register(self, name: str, params: GuardrailParamsBody) -> str: - """Register any guardrail via POST /guardrails and return its id. New - built-ins register with default_on=False and are opted into per request - via the chat body's `guardrails` list, so one guardrail under test never - intercepts unrelated traffic on the shared proxy.""" - return unwrap( + """Register any guardrail via POST /guardrails and return its id, once every + replica can be expected to serve it. New built-ins register with + default_on=False and are opted into per request via the chat body's + `guardrails` list, so one guardrail under test never intercepts unrelated + traffic on the shared proxy. + + /guardrails is a control-plane route and guardrails reach the data plane on + the config reload, so a request naming this guardrail the instant the POST + returns can 404 with "Guardrail not found" on a replica that has not + reloaded. There is no data-plane read that lists guardrails, so unlike + ProxyClient.create_model this settles on the propagation budget alone with + nothing to poll first.""" + guardrail_id = unwrap( self.proxy.transport.post( "/guardrails", headers=self.proxy.transport.master, @@ -188,6 +176,8 @@ class GuardrailsClient: response_type=GuardrailCreateResponse, ) ).guardrail_id + settle_propagation(time.monotonic()) + return guardrail_id def delete_guardrail(self, guardrail_id: str) -> None: _ = self.proxy.transport.delete( diff --git a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py index b6d90b7f6a2..5e9c9f614e5 100644 --- a/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py +++ b/tests/e2e/llm_translation/test_vertex_passthrough_e2e.py @@ -23,11 +23,12 @@ model, spend > 0), correlated by the x-litellm-call-id header. """ import os +import time import pytest from pydantic import BaseModel -from e2e_config import unique_marker +from e2e_config import settle_propagation, unique_marker from e2e_http import NoBody, require_successful_call, unwrap from lifecycle import ResourceManager from models import SpendLogRow @@ -90,7 +91,14 @@ class _ModelDeleteBody(BaseModel): def _add_vertex_passthrough_model( client: PassthroughClient, model_name: str, project: str, credentials: str ) -> str: - return unwrap( + """Register the passthrough deployment and settle before the caller uses it. + + This body carries `use_in_pass_through` and a pinned `model_info.id`, so it + cannot go through ProxyClient.create_model -- but it needs that helper's + propagation settle just the same, or the passthrough call can land on a replica + that has not reloaded yet. + """ + model_id = unwrap( client.proxy.transport.post( "/model/new", headers=client.proxy.transport.master, @@ -108,6 +116,8 @@ def _add_vertex_passthrough_model( response_type=_ModelNewResponse, ) ).model_id + settle_propagation(time.monotonic()) + return model_id def _delete_model(client: PassthroughClient, model_id: str) -> None: diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index 7053bd1dfd2..d76f7b356b2 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -23,7 +23,7 @@ from typing import Callable, Literal import pytest from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError -from e2e_config import POLL_INTERVAL, POLL_TIMEOUT +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT, settle_propagation from proxy_client import ProxyClient from e2e_http import ( URL, @@ -409,6 +409,7 @@ class LoggingClient: ) guardrail_id = response.guardrail_id assert guardrail_id, f"create guardrail returned no id: {response!r}" + settle_propagation(time.monotonic()) return guardrail_id def delete_guardrail(self, guardrail_id: str) -> None: diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 33ec557c339..73453478e5a 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -18,6 +18,7 @@ from dataclasses import dataclass from pydantic import BaseModel, ConfigDict, Field, RootModel +from e2e_config import settle_propagation from e2e_http import Headers, NoBody, Result, Success, UnknownApiError, unwrap from models import KeyGenerateBody, ObjectPermission from proxy_client import ProxyClient @@ -330,7 +331,7 @@ class McpClient: tool-call hook (pre_mcp_call) and blocks a single keyword. The keyword is unique per test, so default_on only ever intercepts this test's own banned tool call on the shared proxy.""" - return unwrap( + guardrail_id = unwrap( self.proxy.transport.post( "/guardrails", headers=self.proxy.transport.master, @@ -345,6 +346,8 @@ class McpClient: response_type=GuardrailCreateResponse, ) ).guardrail_id + settle_propagation(time.monotonic()) + return guardrail_id def delete_guardrail(self, guardrail_id: str) -> None: _ = self.proxy.transport.delete( diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 6c6b948e29c..5050b6fce68 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -70,17 +70,133 @@ from e2e_config import ( POLL_TIMEOUT, PROXY_BASE_URL, REQUEST_TIMEOUT, + settle_propagation, ) from transport import HttpTransport, SplitTransport, Transport RowsPredicate = Callable[[list[SpendLogRow]], bool] +# After /model/new, poll data-plane /v1/models until the model is listed (or fail). +# Bound by MODEL_SERVABLE_TIMEOUT so a stuck reload does not burn the spend +# poll_timeout (120s). Return on first listing; settle_propagation owns the separate +# wait that lets every worker and replica reload before the caller uses the model. +MODEL_SERVABLE_TIMEOUT = 40.0 +MODEL_SERVABLE_DB_SYNC_SECONDS = 0.0 +MODEL_SERVABLE_INTERVAL = 2.0 +# Cap each /v1/models poll so one slow request cannot outlast the remaining budget. +MODEL_SERVABLE_REQUEST_TIMEOUT = 5.0 + + +@dataclass(frozen=True, slots=True) +class Servable: + """The data plane listed the model within the deadline.""" + + +@dataclass(frozen=True, slots=True) +class NotServable: + """The deadline passed without the data plane listing the model. + + `last_result` is the final /v1/models read, so the caller can tell "the proxy + answered but omitted the model" (propagation) from "the read itself failed" + (network/auth) when reporting.""" + + last_result: Result[ModelsListResponse] | None + + +ServableOutcome = Servable | NotServable + + +def await_servable( + list_models: Callable[[float], Result[ModelsListResponse]], + *, + model_name: str, + timeout: float, + interval: float, + request_timeout: float, + db_sync_seconds: float, + now: Callable[[], float], + sleep: Callable[[float], None], +) -> ServableOutcome: + """Poll until `model_name` is listed long enough for every worker to DB-sync. + + First listing must happen within `timeout`. After that, the model must stay + listed continuously for `db_sync_seconds` (any miss resets the continuous + window). `db_sync_seconds=0` returns on the first listing. Each poll's request + timeout is clamped to the remaining budget. Sleeps only min(interval, time left) + so a final deadline-clamped poll is never skipped just because a full interval + does not fit. Clock and sleep are injected.""" + started = now() + first_seen_at: float | None = None + last_result: Result[ModelsListResponse] | None = None + while True: + t = now() + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + remaining = phase_deadline - t + if remaining <= 0: + if ( + last_result is not None + and first_seen_at is not None + and (db_sync_seconds <= 0 or t - first_seen_at >= db_sync_seconds) + ): + return Servable() + return NotServable(last_result=last_result) + + poll_timeout = min(request_timeout, remaining) + last_result = list_models(poll_timeout) + listed = isinstance(last_result, Success) and any( + entry.id == model_name for entry in last_result.data.data + ) + t = now() + if not listed: + first_seen_at = None + elif first_seen_at is None: + if t > started + timeout: + return NotServable(last_result=last_result) + first_seen_at = t + if db_sync_seconds <= 0: + return Servable() + elif t - first_seen_at >= db_sync_seconds: + return Servable() + + phase_deadline = ( + started + timeout if first_seen_at is None else first_seen_at + db_sync_seconds + ) + wait = min(interval, phase_deadline - now()) + if wait > 0: + sleep(wait) + + +def servable_timeout_message( + *, + model_name: str, + timeout: float, + db_sync_seconds: float, + last_result: Result[ModelsListResponse] | None, +) -> str: + last_error = ( + f"; last /v1/models poll did not succeed: {last_result}" + if last_result is not None and not isinstance(last_result, Success) + else "" + ) + return ( + f"model {model_name!r} was created but never became servable on the data " + f"plane within {timeout}s of first listing (plus {db_sync_seconds}s continuous " + f"DB sync) after /model/new (control/data-plane propagation or " + f"STORE_MODEL_IN_DB reload issue){last_error}" + ) + @dataclass(frozen=True, slots=True) class ProxyClient: transport: Transport poll_timeout: float = 120.0 poll_interval: float = 5.0 + model_servable_timeout: float = MODEL_SERVABLE_TIMEOUT + model_servable_db_sync_seconds: float = MODEL_SERVABLE_DB_SYNC_SECONDS + model_servable_interval: float = MODEL_SERVABLE_INTERVAL + model_servable_request_timeout: float = MODEL_SERVABLE_REQUEST_TIMEOUT # ---- keys / customers (satisfies lifecycle.ResourceClient) ---------- @@ -161,13 +277,18 @@ class ProxyClient: """Register a deployment under `model_name` and return its proxy-assigned model_id, once the model is actually servable on the data plane. - /model/new is a control-plane route; in a split control/data-plane - deployment the gateway (data plane, which serves /chat, /ocr, ...) only - picks the new model up on its next DB reload, so a call issued the instant - this returns can race the reload and 400 with "Invalid model name passed". - We therefore poll the data-plane /v1/models until the model appears before - handing back, so callers can invoke it immediately. In the monolithic case - it is already present on the first poll, so this adds one request.""" + /model/new is a control-plane route; the data plane (which serves /chat, + /ocr, ...) only picks the new model up on its next DB reload, so a call + issued the instant this returns can race the reload and 400 with "Invalid + model name passed". We poll the data-plane /v1/models until the model + appears, then settle for the remainder of the propagation budget. + + Both steps are needed, and the second is the one that matters at >1 replica. + The poll proves *a* replica is serving the model; it cannot prove they all + are, because every request opens a fresh connection and a load-balanced + Service routes each one independently -- so the caller's next request + re-rolls and can land on a replica that has not reloaded yet. Waiting out + PROPAGATION_TIMEOUT is what makes the model safe to use anywhere.""" model_id = unwrap( self.transport.post( "/model/new", @@ -180,37 +301,41 @@ class ProxyClient: response_type=ModelNewResponse, ) ).model_id + written_at = time.monotonic() self._await_model_servable(model_name) + settle_propagation(written_at) return model_id def _await_model_servable(self, model_name: str) -> None: - """Block until the data plane lists `model_name`, or fail loudly if it does - not within poll_timeout (a real propagation/config problem, surfaced here - instead of as a downstream "Invalid model name passed").""" - deadline = time.monotonic() + self.poll_timeout - last_result: Result[ModelsListResponse] | None = None - while time.monotonic() < deadline: - last_result = self.transport.get( + """Block until the data plane lists `model_name`, or fail at model_servable_timeout.""" + outcome = await_servable( + lambda poll_timeout: self.transport.get( "/v1/models", headers=self.transport.master, params=NoBody(), response_type=ModelsListResponse, - ) - if isinstance(last_result, Success) and any( - entry.id == model_name for entry in last_result.data.data - ): + timeout=poll_timeout, + ), + model_name=model_name, + timeout=self.model_servable_timeout, + interval=self.model_servable_interval, + request_timeout=self.model_servable_request_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + now=time.monotonic, + sleep=time.sleep, + ) + match outcome: + case Servable(): return - time.sleep(self.poll_interval) - last_error = ( - f"; last /v1/models poll did not succeed: {last_result}" - if last_result is not None and not isinstance(last_result, Success) - else "" - ) - raise AssertionError( - f"model {model_name!r} was created but never became servable on the data " - f"plane within {self.poll_timeout}s of /model/new (control/data-plane " - f"propagation or STORE_MODEL_IN_DB reload issue){last_error}" - ) + case NotServable(last_result=last_result): + raise AssertionError( + servable_timeout_message( + model_name=model_name, + timeout=self.model_servable_timeout, + db_sync_seconds=self.model_servable_db_sync_seconds, + last_result=last_result, + ) + ) def update_model(self, model_id: str, litellm_params: LiteLLMParamsBody) -> None: """Merge `litellm_params` over the deployment `model_id`'s stored params via diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index a6adf83ed1f..27b11befc8e 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -58,6 +58,7 @@ class Transport(Protocol): headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: ... def delete[R: BaseModel]( @@ -136,13 +137,16 @@ class HttpTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: + """`timeout` overrides the transport-wide request_timeout for this call, for + pollers whose own deadline is shorter than it.""" return e2e_http.get( self._url(path), headers=headers, params=params, response_type=response_type, - timeout=self.request_timeout, + timeout=self.request_timeout if timeout is None else timeout, ) def delete[R: BaseModel]( @@ -336,9 +340,14 @@ class SplitTransport: headers: BaseModel, params: BaseModel, response_type: type[R], + timeout: float | None = None, ) -> Result[R]: return self._route(path).get( - path, headers=headers, params=params, response_type=response_type + path, + headers=headers, + params=params, + response_type=response_type, + timeout=timeout, ) def delete[R: 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 50af6465d06..e1e5cc6c532 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1813,6 +1813,190 @@ def _create_unified_batch_id(model_id: str, batch_id: str) -> str: return base64.urlsafe_b64encode(unified_str.encode()).decode().rstrip("=") +def _decode_unified_id(b64_id: str) -> str: + return base64.urlsafe_b64decode(b64_id + "=" * (-len(b64_id) % 4)).decode() + + +def _terminal_batch_record( + unified_batch_uid: str, + raw_input_file_id: str, + raw_output_file_id: str, + raw_error_file_id: str, +): + record = MagicMock() + record.unified_object_id = unified_batch_uid + record.created_by = "owner-user" + record.team_id = "owner-team" + record.status = "cancelled" + record.file_object = json.dumps( + { + "id": "batch-raw-456", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "cancelled", + "created_at": 1234567890, + "input_file_id": raw_input_file_id, + "output_file_id": raw_output_file_id, + "error_file_id": raw_error_file_id, + } + ) + return record + + +@pytest.mark.asyncio +async def test_list_batches_registers_and_returns_unified_output_file_ids(): + """A stored batch blob with raw provider file IDs (e.g. persisted by the cost + poller for a cancelled batch) must be listed with unified managed IDs, and the + output/error files must be registered in the managed file table so GET + /files/{id}/content can route them.""" + from litellm.proxy._types import UserAPIKeyAuth + + unified_batch_uid = _create_unified_batch_id("model-123", "batch-456") + raw_input_file_id = "file-list-in-1" + raw_output_file_id = "file-list-out-1" + raw_error_file_id = "file-list-err-1" + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch" + ).decode() + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [ + _terminal_batch_record( + unified_batch_uid, raw_input_file_id, raw_output_file_id, raw_error_file_id + ) + ] + + input_file_row = MagicMock() + input_file_row.unified_file_id = unified_input_file_id + input_file_row.flat_model_file_ids = [raw_input_file_id] + + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[input_file_row] + ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=None) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=10, + ) + + listed = result["data"][0] + assert listed.id == unified_batch_uid + assert listed.input_file_id == unified_input_file_id + + bulk_lookup = prisma_client.db.litellm_managedfiletable.find_many.await_args + assert set(bulk_lookup.kwargs["where"]["flat_model_file_ids"]["hasSome"]) == { + raw_input_file_id, + raw_output_file_id, + raw_error_file_id, + } + + decoded_output = _decode_unified_id(listed.output_file_id) + assert decoded_output.startswith("litellm_proxy") + assert f"llm_output_file_id,{raw_output_file_id}" in decoded_output + assert "llm_output_file_model_id,model-123" in decoded_output + assert "target_model_names,gpt-5-batch" in decoded_output + + decoded_error = _decode_unified_id(listed.error_file_id) + assert f"llm_output_file_id,{raw_error_file_id}" in decoded_error + + upsert_calls = prisma_client.db.litellm_managedfiletable.upsert.await_args_list + stored_raw_ids = { + c.kwargs["data"]["create"]["flat_model_file_ids"][0] for c in upsert_calls + } + assert stored_raw_ids == {raw_output_file_id, raw_error_file_id} + for c in upsert_calls: + assert c.kwargs["data"]["create"]["created_by"] == "owner-user" + assert c.kwargs["data"]["create"]["team_id"] == "owner-team" + + +@pytest.mark.asyncio +async def test_list_batches_resolves_existing_managed_rows_without_minting(): + """When the raw provider file IDs already have managed file rows, listing must + swap in the existing unified IDs via one bulk lookup for the whole page, with + no per-row queries and no duplicate upserts.""" + from litellm.proxy._types import UserAPIKeyAuth + + unified_input_file_id = base64.urlsafe_b64encode( + b"litellm_proxy:application/octet-stream;unified_id,in-9;target_model_names,gpt-5-batch" + ).decode() + raw_output_file_ids = ["file-list-out-existing-1", "file-list-out-existing-2"] + existing_unified_output_ids = [ + base64.urlsafe_b64encode( + f"litellm_proxy:application/json;unified_id,u-{i};llm_output_file_id,{raw_id}".encode() + ).decode() + for i, raw_id in enumerate(raw_output_file_ids) + ] + + records = [ + _terminal_batch_record( + _create_unified_batch_id("model-123", f"batch-{i}"), + unified_input_file_id, + raw_id, + "", + ) + for i, raw_id in enumerate(raw_output_file_ids) + ] + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = records + + existing_rows = [ + MagicMock(unified_file_id=unified_id, flat_model_file_ids=[raw_id]) + for raw_id, unified_id in zip(raw_output_file_ids, existing_unified_output_ids) + ] + + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=existing_rows + ) + prisma_client.db.litellm_managedfiletable.find_first = AsyncMock() + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=10, + ) + + assert [b.output_file_id for b in result["data"]] == existing_unified_output_ids + prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once() + prisma_client.db.litellm_managedfiletable.find_first.assert_not_awaited() + prisma_client.db.litellm_managedfiletable.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_list_batches_caps_page_size_at_100(): + """The list page size must be capped at 100 rows (matching OpenAI's limit) + even when the caller asks for more, so one request cannot fan out into an + unbounded scan.""" + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="owner-user"), + limit=100000, + ) + + assert ( + prisma_client.db.litellm_managedobjecttable.find_many.await_args.kwargs["take"] + == 101 + ) + assert result["data"] == [] + + @pytest.mark.asyncio async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception(): from litellm.proxy._types import UserAPIKeyAuth diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index c55b66b402b..c29b4c68bb0 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -864,3 +864,178 @@ def test_litellm_project_table_has_timestamp_fields(): fields = LiteLLM_ProjectTable.model_fields assert "created_at" in fields, "LiteLLM_ProjectTable must have created_at field" assert "updated_at" in fields, "LiteLLM_ProjectTable must have updated_at field" + + +@pytest.mark.asyncio +async def test_update_project_invalidates_cached_project_object(monkeypatch): + """ + LIT-3803 regression: auth reads projects cache-first with no freshness check, + so /project/update must evict the cached project. Before the fix, a project + cached with models=[] kept bypassing the new allowlist until the TTL expired. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.auth.auth_checks import get_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + cache = UserApiKeyCache() + + stale_row = MagicMock() + stale_row.model_dump = lambda: {"project_id": project_id, "team_id": None, "models": []} + + mock_prisma = MagicMock() + mock_prisma.jsonify_object = lambda data: data + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=stale_row) + + seeded = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert seeded is not None and seeded.models == [] + + updated_models = ["gemini-2.5-flash-image", "gemini-3.1-flash-lite-preview"] + existing_row = MagicMock(team_id=None, budget_id=None, object_permission_id=None) + updated_row = MagicMock() + updated_row.model_dump = lambda: { + "project_id": project_id, + "team_id": None, + "models": updated_models, + } + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_projecttable.update = AsyncMock(return_value=updated_row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + + await update_project( + data=UpdateProjectRequest(project_id=project_id, models=updated_models), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=updated_row) + refreshed = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert refreshed is not None + assert refreshed.models == updated_models + + +@pytest.mark.asyncio +async def test_delete_project_invalidates_cached_project_object(monkeypatch): + """ + LIT-3803 regression: /project/delete must evict the cached project so auth + stops enforcing (or trusting) a project that no longer exists. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.auth.auth_checks import get_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + cache = UserApiKeyCache() + + row = MagicMock() + row.model_dump = lambda: {"project_id": project_id, "team_id": None, "models": ["gpt-5.5"]} + + mock_prisma = MagicMock() + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=row) + + seeded = await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + assert seeded is not None + + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_projecttable.delete = AsyncMock(return_value=row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) + + await delete_project( + data=DeleteProjectRequest(project_ids=[project_id]), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=None) + assert ( + await get_project_object( + project_id=project_id, prisma_client=mock_prisma, user_api_key_cache=cache + ) + is None + ) + + +@pytest.mark.asyncio +async def test_update_project_succeeds_when_cache_eviction_fails(monkeypatch): + """ + The DB write has already committed when eviction runs, so a cache backend + error must not turn a successful update into a 500; the stale entry is + bounded by the TTL instead. + """ + from unittest.mock import AsyncMock, MagicMock + + project_id = f"project-{uuid.uuid4()}" + failing_cache = MagicMock() + failing_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down")) + + existing_row = MagicMock(team_id=None, budget_id=None, object_permission_id=None) + updated_row = MagicMock() + + mock_prisma = MagicMock() + mock_prisma.jsonify_object = lambda data: data + mock_prisma.db.litellm_projecttable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma.db.litellm_projecttable.update = AsyncMock(return_value=updated_row) + + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma) + monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", failing_cache) + + response = await update_project( + data=UpdateProjectRequest(project_id=project_id, models=["gpt-oss-120b"]), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + assert response is updated_row + failing_cache.async_delete_cache.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_project_eviction_publishes_cross_worker_invalidation(monkeypatch): + """ + Single-worker eviction only fixes the handling worker; the broadcast is what + lets every other worker drop its in-memory copy instead of serving the stale + project until the TTL expires. + """ + from unittest.mock import AsyncMock, patch + + from litellm.proxy.auth.auth_checks import delete_cached_project_object + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + project_id = f"project-{uuid.uuid4()}" + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new=AsyncMock(), + ) as mock_publish: + await delete_cached_project_object( + project_id=project_id, user_api_key_cache=UserApiKeyCache() + ) + + mock_publish.assert_awaited_once_with(cache_key=f"project_id:{project_id}") diff --git a/tests/guardrails_tests/test_zscaler_ai_guard.py b/tests/guardrails_tests/test_zscaler_ai_guard.py index 51c86c15dcb..76e498673b0 100644 --- a/tests/guardrails_tests/test_zscaler_ai_guard.py +++ b/tests/guardrails_tests/test_zscaler_ai_guard.py @@ -396,3 +396,122 @@ async def test_apply_guardrail_block_does_not_log_error(mock_api_call): mock_logger.error.assert_not_called() assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_send_request_uses_default_timeout_when_unconfigured(): + """ + Regression: unconfigured guardrails must keep the historical 5s timeout. + """ + guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1) + + assert guardrail.timeout == 5.0 + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.zscaler_ai_guard.get_async_httpx_client" + ) as mock_get_client: + mock_client = Mock() + mock_client.post = AsyncMock(return_value=Mock(status_code=200)) + mock_get_client.return_value = mock_client + + await guardrail._send_request("http://example.com", {}, {}) + + assert mock_client.post.call_args.kwargs["timeout"] == 5.0 + + +@pytest.mark.asyncio +async def test_send_request_uses_configured_timeout(): + """ + Regression for LIT-5222: a configured timeout must reach the HTTP call. + + Before the fix _send_request passed a module-level constant, so a slow + upstream failed at 5s with `Timeout passed=5` no matter what was configured. + """ + guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=30) + + assert guardrail.timeout == 30 + + with patch( + "litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard.zscaler_ai_guard.get_async_httpx_client" + ) as mock_get_client: + mock_client = Mock() + mock_client.post = AsyncMock(return_value=Mock(status_code=200)) + mock_get_client.return_value = mock_client + + await guardrail._send_request("http://example.com", {}, {}) + + assert mock_client.post.call_args.kwargs["timeout"] == 30 + + +def test_initialize_guardrail_forwards_configured_timeout(): + """ + Regression for LIT-5222: the `timeout` key from config.yaml must survive + initialization. It reaches LitellmParams already, but the initializer used + to drop it before it could reach the guardrail instance. + """ + from litellm.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import ( + initialize_guardrail, + ) + from litellm.types.guardrails import LitellmParams + + litellm_params = LitellmParams( + guardrail="zscaler_ai_guard", + mode="pre_call", + api_key="test_key", + api_base="http://example.com", + policy_id=1, + timeout="30", + ) + + guardrail = initialize_guardrail( + litellm_params, {"guardrail_name": "zscaler-configured-timeout"} + ) + + assert guardrail.timeout == 30.0 + + +def test_config_model_exposes_timeout_to_dashboard(): + """ + The dashboard guardrail form is built from get_config_model(), so the field + has to be declared there for the setting to be reachable outside config.yaml. + """ + config_model = ZscalerAIGuard.get_config_model() + + assert config_model is not None + assert "timeout" in config_model.model_fields + + +@pytest.mark.parametrize("bad_timeout", [0, -1]) +def test_non_positive_timeout_falls_back_to_default(bad_timeout): + """ + Regression: httpx rejects a negative timeout and treats 0 as "fail + immediately", so a non-positive value would break every scan instead of + relaxing the limit the operator was trying to raise. + """ + guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=bad_timeout) + + assert guardrail.timeout == 5.0 + + +def test_update_in_memory_litellm_params_keeps_timeout_resolved(): + """ + Regression: the base implementation copies every LitellmParams attribute + onto the guardrail, so an unset timeout would overwrite the resolved value + with None and silently fall back to the shared client's 600s default. + """ + from litellm.types.guardrails import LitellmParams + + guardrail = ZscalerAIGuard(api_key="test_key", policy_id=1, timeout=30) + assert guardrail.timeout == 30 + + guardrail.update_in_memory_litellm_params( + LitellmParams(guardrail="zscaler_ai_guard", mode="pre_call", api_key="test_key") + ) + assert guardrail.timeout == 5.0 + + guardrail.update_in_memory_litellm_params( + LitellmParams( + guardrail="zscaler_ai_guard", mode="pre_call", api_key="test_key", timeout=45 + ) + ) + assert guardrail.timeout == 45.0 diff --git a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py index 78cb1488533..842a635ecad 100644 --- a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py +++ b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py @@ -11,6 +11,10 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: @@ -64,8 +68,8 @@ async def test_get_agents_allowed_when_not_disabled(): MagicMock(get_agent_list=MagicMock(return_value=[])), ): with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ): result = await get_agents(request=request_mock, user_api_key_dict=user) assert result == [] diff --git a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py index 65448c6281e..c263b8ce381 100644 --- a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py +++ b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py @@ -30,6 +30,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( ) from fastapi import Request from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( _update_metadata_with_tags_in_header, HttpPassThroughEndpointHelpers, @@ -652,3 +653,119 @@ def test_custom_pricing_used_in_cost_calculation(): print(f"Cache-aware cost: {cache_cost}") print("✅ Custom pricing parameters are correctly used in cost calculation") + + +def test_init_kwargs_client_metadata_cannot_spoof_authenticated_identity( + mock_request, mock_user_api_key_dict +): + request = mock_request() + passthrough_payload = PassthroughStandardLoggingPayload( + url="https://test.com", + request_body={}, + ) + authenticated_key = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-user", + key_alias="real-key", + team_alias="Real Team", + user_email="real@example.com", + org_id="real-org", + ) + + result = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=request, + user_api_key_dict=authenticated_key, + passthrough_logging_payload=passthrough_payload, + litellm_call_id="test-call-id", + logging_obj=LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=False, + call_type="test-call-type", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ), + _parsed_body={ + "litellm_metadata": { + "user_api_key_org_id": "victim-org", + "user_api_key_end_user_id": "victim-end-user", + "user_api_key_user_id": "victim-user", + "user_api_key_team_id": "victim-team", + "user_api_key_team_alias": "Victim Team", + "user_api_key_alias": "victim-key", + "user_api_key_user_email": "victim@example.com", + } + }, + ) + + metadata = result["litellm_params"]["metadata"] + assert metadata["user_api_key_user_id"] == "test-user" + assert metadata["user_api_key_team_id"] == "test-team" + assert metadata["user_api_key_team_alias"] == "Real Team" + assert metadata["user_api_key_alias"] == "real-key" + assert metadata["user_api_key_user_email"] == "real@example.com" + assert metadata["user_api_key_org_id"] == "real-org" + assert metadata["user_api_key_end_user_id"] == "test-user" + + +def test_init_kwargs_no_authenticated_identity_field_is_client_settable( + mock_request, mock_user_api_key_dict +): + authenticated_key = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-end-user", + key_alias="real-key", + team_alias="Real Team", + user_email="real@example.com", + org_id="real-org", + organization_alias="Real Org", + project_id="real-project", + project_alias="Real Project", + spend=1.5, + max_budget=10.0, + user_spend=2.5, + user_max_budget=20.0, + team_spend=3.5, + team_max_budget=30.0, + metadata={"real": "auth-metadata"}, + ) + expected = dict( + LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( + user_api_key_dict=authenticated_key + ) + ) + assert len(expected) >= 20 + + spoofed = {key: f"SPOOFED-{key}" for key in expected} + + result = HttpPassThroughEndpointHelpers._init_kwargs_for_pass_through_endpoint( + request=mock_request(), + user_api_key_dict=authenticated_key, + passthrough_logging_payload=PassthroughStandardLoggingPayload( + url="https://test.com", request_body={} + ), + litellm_call_id="test-call-id", + logging_obj=LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=False, + call_type="test-call-type", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="test-function-id", + ), + _parsed_body={"litellm_metadata": dict(spoofed), "metadata": dict(spoofed)}, + ) + + metadata = result["litellm_params"]["metadata"] + survived = { + key: metadata.get(key) + for key in expected + if metadata.get(key) != expected[key] + } + assert survived == {}, f"client-supplied values survived for: {sorted(survived)}" diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index aa734ee22cc..65b70f13a3b 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -38,11 +38,13 @@ async def _turn( tokens: int = 100, spend: float = 0.01, saved: float = 0.02, + tier: "str | None" = None, ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( UPSERT_AUTOROUTER_SESSION_SQL, key, session_id, router, router_type, model, at.isoformat(), tokens, spend, saved, covered, hit, ttl, touched, + tier, ) @@ -195,6 +197,106 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db assert [(row["router_type"], row["sessions"]) for row in matching] == [("complexity", 1), ("quality", 1)] +async def test_tier_turns_count_each_tier_that_served_a_turn(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, tier="simple") + await _turn(db, key, "B", T0 + timedelta(seconds=10), tier="complex") + await _turn(db, key, "A", T0 + timedelta(seconds=20), tier="simple") + + assert (await _row(db, key))["tier_turns"] == {"simple": 2, "complex": 1} + + +async def test_an_untiered_turn_increments_no_tier_counter(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, tier=None) + assert (await _row(db, key))["tier_turns"] == {} + + await _turn(db, key, "A", T0 + timedelta(seconds=10), tier="medium") + await _turn(db, key, "A", T0 + timedelta(seconds=20), tier=None) + row = await _row(db, key) + assert row["tier_turns"] == {"medium": 1} + assert row["turns"] == 3 + + +async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_the_map(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, router_type="complexity", tier="medium") + await _turn(db, key, "A", T0 + timedelta(seconds=10), router_type="quality", tier="2") + await _turn(db, key, "A", T0 + timedelta(seconds=20), router_type="complexity", tier="medium") + + row = await _row(db, key) + assert row["tier_turns"] == {"medium": 2} + assert row["turns"] == 3 + + +async def test_an_out_of_order_turn_still_counts_toward_its_tier(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0 + timedelta(seconds=60), tier="simple") + await _turn(db, key, "A", T0, tier="simple") + + row = await _row(db, key) + assert row["tier_turns"] == {"simple": 2} + assert row["unordered_turns"] == 1 + + +async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier="simple") + await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, tier="simple") + await _turn(db, key, "B", T0 + timedelta(seconds=20), session_id=f"s-{uuid.uuid4()}", router=router, tier="complex") + await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + grouped = next(row for row in rows if row["router_name"] == router) + assert grouped["tier_turns"] == {"simple": 2, "complex": 1} + assert grouped["turns"] == 4 + + +async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn( + db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity", tier="medium" + ) + await _turn( + db, + key, + "A", + T0 + timedelta(seconds=10), + session_id=f"s-{uuid.uuid4()}", + router=router, + router_type="quality", + tier="2", + ) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router} + assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}} + + +async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + grouped = next(row for row in rows if row["router_name"] == router) + assert grouped["tier_turns"] == {} + + async def test_a_miss_that_touched_no_cache_does_not_advance_the_ttl_clock(db): key = f"k-{uuid.uuid4()}" await _turn(db, key, "A", T0, ttl=300) diff --git a/tests/proxy_migration_tests/test_component_image_serves_offline.py b/tests/proxy_migration_tests/test_component_image_serves_offline.py new file mode 100644 index 00000000000..528c2960072 --- /dev/null +++ b/tests/proxy_migration_tests/test_component_image_serves_offline.py @@ -0,0 +1,187 @@ +"""Image-level regression net for the prisma bake in the componentized images. + +The gateway and backend serve requests; they never shell out to the Prisma CLI +(``PrismaManager.setup_database`` is reachable only from ``proxy_cli.py``, which +uvicorn'ing ``gateway.main:app`` bypasses). What they do need is the generated +client's baked query engine, and prisma-python resolves those baked paths +eagerly, with an existence check that propagates EACCES rather than skipping the +candidate. An engine baked under a build-time ``HOME`` is therefore unreadable to +any other runtime uid, and the process dies during startup before +``PRISMA_QUERY_ENGINE_BINARY`` is ever consulted. + +That is what an OpenShift ``restricted-v2`` namespace produces: the image +``USER`` is ignored and an arbitrary uid in GID 0 is assigned instead. The +symptom is not a degraded proxy, it is a proxy that does not serve at all. + +Booting the image the way that deployment does, and requiring it to answer a +request with a live database connection, is what catches the whole class: +a boot as the default uid, or one that reaches the internet, passes even when +the bake is unusable everywhere it actually ships. + +Gated on LITELLM_IMAGE (the tag of the image to exercise) so it is skipped in +the normal unit-test run and exercised only where an image has been built (the +image-scan workflow). Requires a working docker CLI. +""" + +import json +import shutil +import subprocess +import time +import uuid + +import os +import pytest + +IMAGE = os.getenv("LITELLM_IMAGE") +POSTGRES_IMAGE = os.getenv("LITELLM_TEST_POSTGRES_IMAGE", "postgres:16-alpine") +CURL_IMAGE = os.getenv("LITELLM_TEST_CURL_IMAGE", "curlimages/curl:8.11.1") +COMPONENT_PORT = os.getenv("LITELLM_COMPONENT_PORT", "4000") +NON_ROOT_UID = "12345:0" +STARTUP_TIMEOUT_SECONDS = int(os.getenv("LITELLM_COMPONENT_STARTUP_TIMEOUT", "180")) + +pytestmark = [ + pytest.mark.skipif(IMAGE is None, reason="requires a built image (set LITELLM_IMAGE)"), + pytest.mark.skipif(shutil.which("docker") is None, reason="requires the docker CLI"), +] + + +def _docker(*args: str, check: bool = True) -> subprocess.CompletedProcess: + return subprocess.run( + ["docker", *args], capture_output=True, text=True, check=check + ) + + +@pytest.fixture() +def offline_stack(): + """A component container and a fresh Postgres on a network with no egress. + + NON_ROOT_UID is an arbitrary uid in GID 0, the shape OpenShift restricted-v2 + assigns. Postgres and curl are pulled while egress still exists, because the + ``--internal`` network below has none: that is what makes a prisma engine + download (binaries.prisma.sh / npm) fail rather than mask a bake that is not + self-contained. + + The container runs with DISABLE_SCHEMA_UPDATE, since applying the schema is + the migration job's responsibility in this topology and needs the Prisma CLI + these images deliberately omit, and with LITELLM_LOCAL_MODEL_COST_MAP, or the + proxy spends the whole startup budget timing out on a cost-map fetch over the + network it does not have. + + Yields (network_name, component_container). Both are torn down afterwards. + """ + run_id = f"componentserve-{uuid.uuid4().hex[:8]}" + network = f"{run_id}-net" + pg = f"{run_id}-pg" + component = f"{run_id}-app" + + _docker("pull", "--quiet", POSTGRES_IMAGE) + _docker("pull", "--quiet", CURL_IMAGE) + _docker("network", "create", "--internal", network) + try: + _docker( + "run", "-d", "--name", pg, "--network", network, + "-e", "POSTGRES_PASSWORD=pw", "-e", "POSTGRES_DB=litellm", + POSTGRES_IMAGE, + ) + _wait_until_postgres_ready(pg) + assert IMAGE is not None + _docker( + "run", "-d", "--name", component, "--network", network, + "--user", NON_ROOT_UID, + "-e", f"DATABASE_URL=postgresql://postgres:pw@{pg}:5432/litellm", + "-e", "LITELLM_MASTER_KEY=sk-component-serve-test", + "-e", "DISABLE_SCHEMA_UPDATE=true", + "-e", "LITELLM_LOCAL_MODEL_COST_MAP=True", + IMAGE, + ) + yield network, component + finally: + _docker("logs", component, check=False) + _docker("rm", "-f", component, check=False) + _docker("rm", "-f", pg, check=False) + _docker("network", "rm", network, check=False) + + +def _wait_until_postgres_ready(pg: str, attempts: int = 60) -> None: + for _ in range(attempts): + running = _docker( + "ps", "--filter", f"name={pg}", "--filter", "status=running", + "--format", "{{.Names}}", check=False, + ).stdout + if pg not in running: + logs = _docker("logs", pg, check=False) + pytest.fail(f"postgres container is not running:\n{logs.stdout}\n{logs.stderr}") + ready = _docker( + "exec", pg, "pg_isready", "-U", "postgres", "-d", "litellm", check=False + ) + if ready.returncode == 0: + return + time.sleep(1) + pytest.fail(f"postgres never became ready after {attempts}s") + + +def _container_logs(container: str) -> str: + logs = _docker("logs", container, check=False) + return f"stdout:\n{logs.stdout}\nstderr:\n{logs.stderr}" + + +def _is_running(container: str) -> bool: + return bool( + _docker( + "ps", "--filter", f"name={container}", "--filter", "status=running", + "--format", "{{.Names}}", check=False, + ).stdout.strip() + ) + + +def _readiness(network: str, component: str) -> subprocess.CompletedProcess: + """Ask the component for its readiness, from a peer on the same egress-less network.""" + return _docker( + "run", "--rm", "--network", network, CURL_IMAGE, + "--silent", "--max-time", "10", + f"http://{component}:{COMPONENT_PORT}/health/readiness", + check=False, + ) + + +def test_component_serves_offline_as_non_root_uid(offline_stack): + """The component answers a request with a live DB connection, offline, as an arbitrary uid. + + On the pre-fix image this never gets a response: the engine baked under + /home/nonroot (mode 0700, owned by uid 65532) raises + ``PermissionError: .../query-engine-linux-...`` out of pathlib and uvicorn + reports ``Application startup failed. Exiting.``. A bake at the fixed, + world-readable /opt/prisma is what lets any uid start the client. + + `db: connected` is the load-bearing part of the assertion: it means the + query engine binary was found, executed, and reached Postgres. A liveness + probe alone would pass on an image whose engine never resolved. + """ + network, component = offline_stack + + deadline = time.time() + STARTUP_TIMEOUT_SECONDS + probe = None + while time.time() < deadline: + if not _is_running(component): + pytest.fail( + f"the component exited during startup as uid {NON_ROOT_UID} with no egress. " + "The prisma bake is not readable to a uid other than the one that built it, " + "so the proxy does not serve at all.\n" + f"{_container_logs(component)}" + ) + probe = _readiness(network, component) + if probe.returncode == 0 and probe.stdout.strip(): + break + time.sleep(2) + + assert probe is not None and probe.returncode == 0 and probe.stdout.strip(), ( + f"/health/readiness never answered within {STARTUP_TIMEOUT_SECONDS}s as uid " + f"{NON_ROOT_UID} with no egress.\n{_container_logs(component)}" + ) + + payload = json.loads(probe.stdout) + assert payload.get("db") == "connected", ( + f"the component answered but its database is {payload.get('db')!r}, so the baked " + f"query engine did not resolve as uid {NON_ROOT_UID}.\nresponse: {probe.stdout}\n" + f"{_container_logs(component)}" + ) diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index bcc70fae67c..0655763d41b 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -15,6 +15,7 @@ from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload from litellm.types.utils import StandardLoggingPayload from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo +from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS @pytest.fixture @@ -1816,6 +1817,7 @@ def test_init_auto_router_deployment_success(mock_auto_router, model_list): default_model="gpt-5-mini", embedding_model="text-embedding-3-small", litellm_router_instance=router, + max_input_chars=DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ) # Verify the auto-router was added to the router's auto_routers dict diff --git a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py index 44d2803a260..06191d1a370 100644 --- a/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py +++ b/tests/test_litellm/a2a_protocol/test_a2a_exception_mapping_utils.py @@ -58,6 +58,32 @@ async def test_localhost_retry_reuses_stashed_httpx_client(): assert mock_create.await_count == 1 +@pytest.mark.asyncio +async def test_localhost_retry_carries_the_agents_call_context_onto_the_new_client(): + """Per-caller headers ride on the call context now, not on the shared httpx client, + so a retry that drops the context would replay the request unauthenticated.""" + stashed_context = object() + a2a_client = MagicMock() + a2a_client._litellm_httpx_client = object() + a2a_client._litellm_call_context = stashed_context + new_client = MagicMock() + + with ( + patch.object(emu, "A2A_SDK_AVAILABLE", True), + patch.object(emu, "set_agent_card_url"), + patch.object(emu, "ClientConfig", side_effect=lambda **_: MagicMock()), + patch.object(emu, "create_client", new=AsyncMock(return_value=new_client)), + ): + result = await emu.handle_a2a_localhost_retry( + error=_localhost_error(), + agent_card=MagicMock(), + a2a_client=a2a_client, + is_streaming=False, + ) + + assert result._litellm_call_context is stashed_context + + @pytest.mark.asyncio async def test_localhost_retry_raises_when_no_stashed_client(): """An externally-supplied client has no LiteLLM httpx handle; the retry must fail diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index c5e48620fde..08f6b9f25bb 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -1,13 +1,26 @@ """Tests for litellm/a2a_protocol/main.py non-streaming send behavior.""" +import httpx import pytest pytest.importorskip("a2a.compat.v0_3.conversions") from a2a.compat.v0_3 import conversions as _conv -from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest +from a2a.compat.v0_3.types import ( + MessageSendParams, + SendMessageRequest, + SendStreamingMessageRequest, +) -from litellm.a2a_protocol.main import _send_message +import litellm +from litellm.a2a_protocol.main import _send_message, _stream_messages, create_a2a_client +from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, + httpxSpecialProvider, +) def _request() -> SendMessageRequest: @@ -40,7 +53,8 @@ class _FakeClient: def __init__(self, *events): self._events = events - async def send_message(self, _pb_request): + async def send_message(self, _pb_request, context=None): + self.context = context for event in self._events: yield event @@ -133,3 +147,198 @@ def test_streaming_logging_obj_carries_call_type_into_model_call_details(): ) assert logging_obj.model_call_details["call_type"] == "asend_message_streaming" + + +_AGENT_CARD = { + "protocolVersion": "0.3.0", + "name": "recording-agent", + "url": "http://127.0.0.1:9/", + "preferredTransport": "JSONRPC", + "version": "1.0.0", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text/plain"], + "defaultOutputModes": ["text/plain"], + "skills": [], +} + +_RPC_REPLY = { + "jsonrpc": "2.0", + "id": "reply", + "result": { + "messageId": "reply-1", + "role": "agent", + "parts": [{"kind": "text", "text": "pong"}], + "kind": "message", + }, +} + +_AGENT_A_HEADERS = {"x-agent-token": "token-for-a", "x-tenant": "tenant-a"} +_AGENT_B_HEADERS = {"x-agent-token": "token-for-b", "x-tenant": "tenant-b"} + + +class _RequestRecorder: + """Records the headers httpx put on the wire, per outbound request.""" + + def __init__(self): + self.card_requests = [] + self.rpc_requests = [] + self.client = None + + def __call__(self, request: httpx.Request) -> httpx.Response: + headers = {k.lower(): v for k, v in request.headers.items()} + if request.method == "GET": + self.card_requests.append(headers) + return httpx.Response(200, json=_AGENT_CARD) + self.rpc_requests.append(headers) + return httpx.Response(200, json=_RPC_REPLY) + + +def _a2a_client_cache_key(timeout: float) -> str: + return "async_httpx_client" + f"timeout_{timeout}" + httpxSpecialProvider.A2AProvider + + +async def _seed_shared_a2a_client() -> _RequestRecorder: + """Put the one A2A client the cache will hand out behind a mock transport. + + Seeding has to happen on the test's own event loop, because the client cache keys on + it. The injected client is a real httpx.AsyncClient, so the merge of per-request + headers over client defaults, which is what these tests are about, stays real. + """ + recorder = _RequestRecorder() + handler = AsyncHTTPHandler(timeout=DEFAULT_A2A_AGENT_TIMEOUT) + owned_client = handler.client + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(recorder)) + await owned_client.aclose() + + litellm.in_memory_llm_clients_cache.set_cache(key=_a2a_client_cache_key(DEFAULT_A2A_AGENT_TIMEOUT), value=handler) + seeded = get_async_httpx_client( + llm_provider=httpxSpecialProvider.A2AProvider, + params={"timeout": DEFAULT_A2A_AGENT_TIMEOUT}, + ) + assert seeded is handler, "cache key drifted from get_async_httpx_client; these tests would test nothing" + + recorder.client = handler.client + return recorder + + +@pytest.fixture +def isolated_client_cache(): + previous = getattr(litellm, "in_memory_llm_clients_cache", None) + litellm.in_memory_llm_clients_cache = LLMClientCache() + yield litellm.in_memory_llm_clients_cache + litellm.in_memory_llm_clients_cache = previous + + +def _send_request(request_id): + return SendMessageRequest( + id=request_id, + params=MessageSendParams( + message={"messageId": request_id, "role": "user", "parts": [{"kind": "text", "text": "hi"}]} + ), + ) + + +@pytest.mark.asyncio +async def test_extra_headers_never_land_on_the_shared_cached_client(isolated_client_cache): + """get_async_httpx_client hands back a process-wide shared client, so a caller's + headers written onto it would outlive the request that supplied them.""" + recorder = await _seed_shared_a2a_client() + + await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS) + + assert "x-agent-token" not in recorder.client.headers + assert "x-tenant" not in recorder.client.headers + + +@pytest.mark.asyncio +async def test_callers_with_different_headers_reuse_one_pooled_client(isolated_client_cache): + """Headers must not segregate the connection pool. Every A2A caller on one timeout + shares one cached client, so header sets cannot multiply cached clients.""" + recorder = await _seed_shared_a2a_client() + + client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS) + client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS) + client_none = await create_a2a_client(base_url="http://127.0.0.1:9") + + assert client_a._litellm_httpx_client is recorder.client + assert client_b._litellm_httpx_client is recorder.client + assert client_none._litellm_httpx_client is recorder.client + + cached = [key for key in isolated_client_cache.cache_dict if "a2a_provider" in key] + assert len(cached) == 1, f"expected one pooled A2A client, cached: {cached}" + + +@pytest.mark.parametrize("order", [("a", "b", "none"), ("b", "none", "a"), ("none", "a", "b")]) +@pytest.mark.asyncio +async def test_each_caller_sends_only_its_own_headers(order, isolated_client_cache): + """Whatever order callers arrive in, each request carries that caller's headers and + no other caller's, and a caller with no extra_headers sends none.""" + recorder = await _seed_shared_a2a_client() + headers_by_caller = {"a": _AGENT_A_HEADERS, "b": _AGENT_B_HEADERS, "none": None} + + for caller in order: + a2a_client = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=headers_by_caller[caller]) + await _send_message(a2a_client, _send_request(caller)) + + received = dict(zip(order, recorder.rpc_requests, strict=True)) + assert received["a"]["x-agent-token"] == "token-for-a" + assert received["a"]["x-tenant"] == "tenant-a" + assert received["b"]["x-agent-token"] == "token-for-b" + assert received["b"]["x-tenant"] == "tenant-b" + assert "x-agent-token" not in received["none"] + assert "x-tenant" not in received["none"] + + +@pytest.mark.asyncio +async def test_streaming_send_carries_only_its_own_caller_headers(isolated_client_cache): + """The streaming path shares the same pooled client, so it needs the same guard.""" + recorder = await _seed_shared_a2a_client() + + client_a = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS, streaming=True) + await _send_message(client_a, _send_request("a")) + + client_b = await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_B_HEADERS, streaming=True) + streaming_request = SendStreamingMessageRequest( + id="b", + params=MessageSendParams(message={"messageId": "b", "role": "user", "parts": [{"kind": "text", "text": "hi"}]}), + ) + async for _ in _stream_messages(client_b, streaming_request): + pass + + received = dict(zip(("a", "b"), recorder.rpc_requests, strict=True)) + assert received["a"]["x-agent-token"] == "token-for-a" + assert received["b"]["x-agent-token"] == "token-for-b" + assert received["b"]["x-tenant"] == "tenant-b" + + +@pytest.mark.asyncio +async def test_agent_card_fetch_carries_the_callers_headers(isolated_client_cache): + """Agent cards can sit behind the same auth as the agent, so the card fetch must stay + authenticated once the headers stop living on the client.""" + recorder = await _seed_shared_a2a_client() + + await create_a2a_client(base_url="http://127.0.0.1:9", extra_headers=_AGENT_A_HEADERS) + + assert recorder.card_requests, "no agent card request was made" + assert recorder.card_requests[-1]["x-agent-token"] == "token-for-a" + + +@pytest.mark.asyncio +async def test_the_pooled_a2a_client_arrives_with_cookie_persistence_disabled(isolated_client_cache): + """create_a2a_client takes its client from the shared builder rather than building one, + and the builder is what refuses to persist cookies. This pins the join between those + two facts, so the A2A path cannot quietly start acquiring a client that keeps a jar. + + test_callers_with_different_headers_reuse_one_pooled_client pins the other half, that + create_a2a_client hands back exactly this cached client.""" + handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.A2AProvider, + params={"timeout": DEFAULT_A2A_AGENT_TIMEOUT}, + ) + request = httpx.Request("GET", "https://agent-a.example.com/") + handler.client.cookies.extract_cookies( + httpx.Response(200, headers={"set-cookie": "SESSION=only-agent-a-may-hold-this"}, request=request) + ) + + assert dict(handler.client.cookies) == {}, "the pooled A2A client kept an upstream's cookie" + await handler.close() diff --git a/tests/test_litellm/conftest.py b/tests/test_litellm/conftest.py index c0993051d33..0dc8f56f3ce 100644 --- a/tests/test_litellm/conftest.py +++ b/tests/test_litellm/conftest.py @@ -19,6 +19,7 @@ sys.path.insert( import asyncio import litellm +from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module from litellm._logging import ALL_LOGGERS from litellm.litellm_core_utils.prompt_templates import ( @@ -244,6 +245,8 @@ def isolate_litellm_state(): for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items() } + original_live_routers = set(litellm_router_module._live_routers) + # Store LiteLLM logger state. Some tests reconfigure handlers/propagation for # JSON logging and do not restore them, which breaks later caplog-based tests. logger_state = {} @@ -313,6 +316,11 @@ def isolate_litellm_state(): litellm_utils_module._runtime_registered_model_cost.clear() litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost) + for _router in tuple(litellm_router_module._live_routers): + litellm_router_module._live_routers.discard(_router) + for _router in original_live_routers: + litellm_router_module._live_routers.add(_router) + # Restore logger configuration mutated by logging-focused tests. for logger in ALL_LOGGERS: original_logger_state = logger_state.get(logger.name) 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 c18db9fe40f..63646ae53f8 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -8,6 +8,7 @@ async_post_call_success_hook when processing completed batch responses. import asyncio import base64 import json +import logging import pytest from typing import Optional @@ -189,6 +190,47 @@ async def test_get_user_created_file_ids_remaps_stored_raw_provider_id_to_unifie assert files[0].purpose == raw_provider_object.purpose +@pytest.mark.asyncio +async def test_parse_managed_file_object_warning_omits_rejected_values(caplog): + from litellm_enterprise.proxy.hooks.managed_files import ( + _parse_managed_file_object, + ) + + with caplog.at_level(logging.WARNING): + parsed = _parse_managed_file_object( + {"id": "file-corrupt", "object": "file", "filename": "confidential.jsonl"}, + "unified-corrupt", + ) + + assert parsed is None + assert "unified-corrupt" in caplog.text + assert "bytes" in caplog.text + assert "confidential.jsonl" not in caplog.text + + +@pytest.mark.asyncio +async def test_get_user_created_file_ids_skips_unparseable_rows(): + managed_files = _make_managed_files_instance() + managed_files.prisma_client.db.litellm_managedfiletable.find_many = AsyncMock( + return_value=[ + MagicMock( + file_object={"id": "file-corrupt", "object": "file"}, + unified_file_id="unified-corrupt", + ), + MagicMock( + file_object=_make_file_object().model_dump(), + unified_file_id="unified-valid", + ), + ] + ) + + files = await managed_files.get_user_created_file_ids( + _make_user_api_key_dict(), ["file-output-abc"] + ) + + assert [file.id for file in files] == ["unified-valid"] + + @pytest.mark.asyncio async def test_should_fallback_when_no_router(): """ diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 41c02501acc..82b074220fa 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -356,6 +356,7 @@ def _mcp_payload(**overrides): "arguments": {"city": "Paris"}, "result": {"temp_c": 21}, "mcp_server_name": "weather-mcp", + "mcp_server_resource": "https://weather.example.com", "mcp_session_id": "sess-abc123", }, }, @@ -452,6 +453,8 @@ def test_mcp_tool_call_failure_marks_error(): assert span.name == "tools/call get_weather" assert span.status.status_code is StatusCode.ERROR assert span.attributes["error.type"] == "MCPError" + assert span.attributes["rpc.system"] == "jsonrpc" + assert span.attributes["server.address"] == "weather.example.com" def test_mcp_tool_call_deduped_on_repeat(): @@ -531,6 +534,82 @@ _MCP_SPAN_CASES = [ ] +def test_mcp_tool_call_names_its_rpc_system_and_upstream(): + """A tool-call span names the RPC system, and always alongside the upstream it called. + + A CLIENT span holding none of the ``rpc.*``/``http.*``/``db.*``/``messaging.*`` + families is unclassifiable, so a backend deriving a span type from them has nothing + to derive from: Elastic APM indexed these spans as ``span.type=unknown`` with no + ``span.subtype`` at all, and its span-links API then rejected the whole trace with + ``Missing required fields (span.subtype)``. + + The two assertions are one invariant, not two. Naming the RPC system makes a + consumer treat the span as a downstream dependency and key that dependency off + ``server.address``/``server.port``; emitting the first without the second names the + dependency ``:0``, which is worse than leaving the span unclassified. + """ + logger, exporter = _logger() + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": _mcp_payload()}, None, None, None + ) + ) + (span,) = exporter.get_finished_spans() + assert span.attributes["rpc.system"] == "jsonrpc" + assert span.attributes["server.address"] == "weather.example.com" + assert span.attributes["server.port"] == 443 + + +@pytest.mark.parametrize( + "resource", + [None, "mcp://weather.example.com", "ws://weather.example.com"], + ids=["no-resource", "scheme-with-no-default-port", "ws-scheme"], +) +def test_mcp_tool_call_omits_rpc_system_without_a_complete_upstream(resource): + """A tool call drops the RPC system unless the full destination resolved. + + ``mcp_server_resource`` is absent whenever the tool name resolves to no registered + server, and it is ``None`` for a transport with no host to log at all (stdio). A + host-bearing scheme outside the HTTP(S) default-port map resolves an address but no + port, and the ``url`` field is not scheme-validated, so that state is reachable from + config. Each case names the dependency ``:0`` or ``host:0`` if ``rpc.system`` ships + on its own, so the pairing is enforced here rather than left to the extractors + happening to agree. + """ + logger, exporter = _logger() + payload = _mcp_payload() + if resource is None: + del payload["metadata"]["mcp_tool_call_metadata"]["mcp_server_resource"] + else: + payload["metadata"]["mcp_tool_call_metadata"]["mcp_server_resource"] = resource + asyncio.run( + logger.async_log_success_event({"standard_logging_object": payload}, None, None, None) + ) + (span,) = exporter.get_finished_spans() + assert "rpc.system" not in span.attributes + assert "server.port" not in span.attributes + assert span.attributes["mcp.method.name"] == "tools/call" + + +def test_mcp_list_tools_omits_rpc_system_without_an_upstream(): + """The discovery span carries no upstream identity, so it must not claim to be RPC. + + ``tools/list`` reaches the callbacks with no ``mcp_tool_call_metadata``, so there is + no ``server.address`` to attach and a listing can span several upstreams anyway. + Naming ``rpc.system`` here would buy a ``span.subtype`` at the cost of a bogus ``:0`` + dependency node in every consumer that aggregates on it. + """ + logger, exporter = _logger() + asyncio.run( + logger.async_log_success_event( + {"standard_logging_object": _mcp_list_payload()}, None, None, None + ) + ) + (span,) = exporter.get_finished_spans() + assert "rpc.system" not in span.attributes + assert "server.address" not in span.attributes + + @pytest.mark.parametrize("make_payload, span_name", _MCP_SPAN_CASES) def test_mcp_span_nests_under_transport_without_propagated_context( make_payload, span_name @@ -1116,8 +1195,13 @@ def test_async_post_call_failure_hook_skips_a_transport_that_already_answered(): def test_record_error_attributes_on_span_decorates_without_ending(): """PATH A: a failure that dies before any LLM-call span (malformed body, validation) is stamped onto the instrumentor-owned SERVER span. The method must - not end the span or emit a duplicate exception event, and must pin error.code - to the real response status (not the exception's own code).""" + not end the span, and must pin error.code to the real response status (not the + exception's own code). + + LIT-4780: the instrumentor never sees the exception (the proxy handler turns it + into a JSONResponse), so nothing else marks the span as failed; the status and + the exception event have to come from here or the trace shows the error message + on an otherwise successful-looking request.""" logger, exporter = _logger() server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) logger.record_error_attributes_on_span(server, _proxy_exc("Invalid JSON body", 400), 422) @@ -1127,7 +1211,31 @@ def test_record_error_attributes_on_span_decorates_without_ending(): assert span.attributes["error.type"] == "ProxyException" assert span.attributes["error.message"] == "Invalid JSON body" assert span.attributes["litellm.provider.error.code"] == "422" - assert all(e.name != "exception" for e in span.events) + assert span.status.status_code is StatusCode.ERROR + assert [e.name for e in span.events] == ["exception"] + + +def test_record_error_attributes_on_span_does_not_duplicate_an_already_stamped_error(): + """A failure that already went through ``async_post_call_failure_hook`` reaches + the exception handler too; the second stamp must keep one exception event while + still repinning error.code to the real response status.""" + from litellm.proxy._types import UserAPIKeyAuth + + logger, exporter = _logger() + server = logger._emitter.start_span(SpanRole.PROXY_REQUEST, LITELLM_PROXY_REQUEST_SPAN_NAME) + set_request_root_span(server) + exc = _proxy_exc("Authentication Error, invalid key", 401) + asyncio.run( + logger.async_post_call_failure_hook( + request_data={}, original_exception=exc, user_api_key_dict=UserAPIKeyAuth() + ) + ) + logger.record_error_attributes_on_span(server, exc, 400) + server.end() + (span,) = exporter.get_finished_spans() + assert [e.name for e in span.events] == ["exception"] + assert span.attributes["litellm.provider.error.code"] == "400" + assert span.status.status_code is StatusCode.ERROR def test_record_error_attributes_on_span_ignores_below_400_and_missing_span(): diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index 612ac1e5113..19d0cfc0b18 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -24,7 +24,11 @@ from litellm.integrations.otel import ( resolve_provider, ) from litellm.integrations.otel.model import spans as spans_mod -from litellm.integrations.otel.model.payloads import LLMCallSpanData, RequestIdentity +from litellm.integrations.otel.model.payloads import ( + LLMCallSpanData, + RequestIdentity, + _upstream_address_port, +) from litellm.integrations.otel.model.spans import ( SPAN_REGISTRY, LiteLLMSpanKind, @@ -177,6 +181,7 @@ def test_mcp_attribute_vocabulary_is_complete(): "mcp.resource.uri", "jsonrpc.request.id", "jsonrpc.protocol.version", + "rpc.system", "rpc.response.status_code", "gen_ai.operation.name", "gen_ai.tool.name", @@ -808,3 +813,34 @@ def test_promoted_baggage_is_bounded_allowlist(): # http.* is never a promoted key assert HTTP.ROUTE not in promoted assert HTTP.REQUEST_METHOD not in promoted + + +@pytest.mark.parametrize( + "resource, expected", + [ + ("https://weather.example.com", ("weather.example.com", 443)), + ("http://weather.example.com", ("weather.example.com", 80)), + ("https://weather.example.com:8443", ("weather.example.com", 8443)), + ("mcp://weather.example.com", ("weather.example.com", None)), + ("http://::1:8080", (None, None)), + ("http://fe80::1%25eth0:80", (None, None)), + (None, (None, None)), + ("", (None, None)), + ], + ids=[ + "https-default", + "http-default", + "explicit-port", + "no-default-port", + "ipv6-unbracketed", + "ipv6-zone-scoped", + "none", + "empty", + ], +) +def test_upstream_address_port(resource, expected): + """The redacted MCP origin resolves to the address and port a consumer names its + dependency from. A scheme outside the default-port map yields no port, and an IPv6 + origin yields nothing at all because the redactor rebuilds it without its brackets; + both are why the mapper gates ``rpc.system`` on the complete pair.""" + assert _upstream_address_port(resource) == expected diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 55e462c82b8..f48f5cb1784 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -296,3 +296,132 @@ async def test_azure_sentinel_audit_stream_name_from_env_var(monkeypatch): ) assert explicit_logger.audit_stream_name == "Custom-LiteLLM-Explicit" + + +def _build_logger(**overrides): + kwargs = { + "dcr_immutable_id": "dcr-test123456789", + "endpoint": "https://test-dce.eastus-1.ingest.monitor.azure.com", + "tenant_id": "test-tenant-id", + "client_id": "test-client-id", + "client_secret": "test-client-secret", + **overrides, + } + with patch("asyncio.create_task", side_effect=_close_periodic_flush_task): + return AzureSentinelLogger(**kwargs) + + +@pytest.fixture +def _no_authority_host_env(monkeypatch): + monkeypatch.delenv("AZURE_SENTINEL_AUTHORITY_HOST", raising=False) + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + + +@pytest.mark.parametrize( + "authority_host, expected_authority, expected_scope", + [ + (None, "https://login.microsoftonline.com", "https://monitor.azure.com/.default"), + ("https://login.microsoftonline.us", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("https://login.microsoftonline.us/", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("login.microsoftonline.us", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("https://adfs.contoso.example", "https://adfs.contoso.example", "https://monitor.azure.com/.default"), + ], +) +def test_azure_sentinel_resolves_authority_host_and_audience_together( + _no_authority_host_env, authority_host, expected_authority, expected_scope +): + """Both the Entra authority and the Azure Monitor audience must follow the configured cloud. + + Moving only the authority leaves a sovereign deployment asking sovereign Entra for the + commercial audience, which the sovereign ingestion endpoint rejects. + """ + logger = _build_logger(**({} if authority_host is None else {"authority_host": authority_host})) + + assert logger.authority_host == expected_authority + assert logger.oauth_scope == expected_scope + + +def test_azure_sentinel_authority_host_from_env_var(_no_authority_host_env, monkeypatch): + """AZURE_AUTHORITY_HOST is the documented setting and the string callback constructs the logger + with no arguments, so the env var alone has to move both values.""" + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger() + + assert logger.authority_host == "https://login.microsoftonline.us" + assert logger.oauth_scope == "https://monitor.azure.us/.default" + + +@pytest.mark.asyncio +async def test_azure_sentinel_token_request_uses_sovereign_authority_and_audience(_no_authority_host_env): + """The resolved values must reach the wire, not just the instance attributes.""" + logger = _build_logger(authority_host="https://login.microsoftonline.us") + logger.log_queue.append( + StandardLoggingPayload( + id="test_id", + call_type="completion", + model="gpt-3.5-turbo", + status="success", + messages=[{"role": "user", "content": "Hello"}], + response={"choices": [{"message": {"content": "Hi"}}]}, + ) + ) + + mock_token_response = MagicMock() + mock_token_response.status_code = 200 + mock_token_response.json = MagicMock(return_value={"access_token": "test-bearer-token", "expires_in": 3600}) + mock_token_response.text = "Success" + mock_api_response = MagicMock() + mock_api_response.status_code = 204 + mock_api_response.text = "Success" + + async def mock_post(*args, **kwargs): + if "oauth2/v2.0/token" in kwargs.get("url", ""): + return mock_token_response + return mock_api_response + + logger.async_httpx_client.post = AsyncMock(side_effect=mock_post) + + await logger.async_send_batch() + + token_calls = [ + call for call in logger.async_httpx_client.post.call_args_list if "oauth2/v2.0/token" in call.kwargs["url"] + ] + assert len(token_calls) == 1 + assert token_calls[0].kwargs["url"] == "https://login.microsoftonline.us/test-tenant-id/oauth2/v2.0/token" + assert token_calls[0].kwargs["data"]["scope"] == "https://monitor.azure.us/.default" + + +def test_azure_sentinel_authority_host_prefers_the_sentinel_scoped_env_var(_no_authority_host_env, monkeypatch): + """AZURE_AUTHORITY_HOST is shared with Azure OpenAI and the azure_storage callback, so a deployment + whose Sentinel workspace lives in a different cloud than the rest of its Azure resources needs a + Sentinel-scoped override. This mirrors how tenant, client id and secret already resolve.""" + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.com") + monkeypatch.setenv("AZURE_SENTINEL_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger() + + assert logger.authority_host == "https://login.microsoftonline.us" + assert logger.oauth_scope == "https://monitor.azure.us/.default" + + +def test_azure_sentinel_falls_back_to_the_shared_authority_host(_no_authority_host_env, monkeypatch): + """With no Sentinel-scoped override the shared variable still applies, which is the behavior + shipped in the original fix.""" + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger() + + assert logger.authority_host == "https://login.microsoftonline.us" + assert logger.oauth_scope == "https://monitor.azure.us/.default" + + +def test_azure_sentinel_authority_host_argument_outranks_the_scoped_env_var(_no_authority_host_env, monkeypatch): + """An explicit constructor argument is the most specific source and has to win, otherwise a + deployment that exports the scoped variable silently overrides an SDK caller.""" + monkeypatch.setenv("AZURE_SENTINEL_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger(authority_host="https://login.microsoftonline.com") + + assert logger.authority_host == "https://login.microsoftonline.com" + assert logger.oauth_scope == "https://monitor.azure.com/.default" diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py index b6ff3b70a4d..f39f41a6d12 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_interception_handler.py @@ -221,14 +221,97 @@ async def test_execute_search_passes_selected_search_tool_litellm_params(monkeyp kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}}, ) - mock_asearch.assert_awaited_once_with( - query="what is litellm", - search_provider="tavily", - api_key="fake-ui-key", - api_base="https://api.tavily.com", - timeout=10.0, - max_retries=2, + forwarded_kwargs = mock_asearch.await_args.kwargs + assert forwarded_kwargs["query"] == "what is litellm" + assert forwarded_kwargs["search_provider"] == "tavily" + assert forwarded_kwargs["api_key"] == "fake-ui-key" + assert forwarded_kwargs["api_base"] == "https://api.tavily.com" + assert forwarded_kwargs["timeout"] == 10.0 + assert forwarded_kwargs["max_retries"] == 2 + + +@pytest.mark.asyncio +async def test_execute_search_attributes_spend_to_the_calling_key(monkeypatch): + """An intercepted search is billed and logged against the key that made the LLM request. + + Without the forwarded attribution metadata the proxy's spend hook skips the search + entirely, so its provider cost never reaches SpendLogs or any budget. + """ + import litellm + from litellm.proxy import proxy_server + from litellm.proxy.hooks.proxy_track_cost_callback import _should_track_cost_callback + + logger = WebSearchInterceptionLogger( + enabled_providers=["bedrock"], + search_tool_name="perplexity-sonar-pro", ) + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "perplexity-sonar-pro", + "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"}, + } + ] + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + user_api_key_auth = UserAPIKeyAuth( + api_key="hashed-sk-1234", + key_alias="alice-key", + user_id="user-alice", + org_id="org-1", + ) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search( + "what is litellm", + kwargs={"litellm_params": {"metadata": {"user_api_key_auth": user_api_key_auth}}}, + ) + + forwarded_metadata = mock_asearch.await_args.kwargs["litellm_metadata"] + assert forwarded_metadata["user_api_key"] == "hashed-sk-1234" + assert forwarded_metadata["user_api_key_hash"] == "hashed-sk-1234" + assert forwarded_metadata["user_api_key_alias"] == "alice-key" + assert forwarded_metadata["user_api_key_user_id"] == "user-alice" + assert forwarded_metadata["user_api_key_org_id"] == "org-1" + assert forwarded_metadata["model_group"] == "perplexity-sonar-pro" + assert ( + _should_track_cost_callback( + user_api_key=forwarded_metadata["user_api_key"], + user_id=forwarded_metadata["user_api_key_user_id"], + team_id=forwarded_metadata["user_api_key_team_id"], + end_user_id=None, + call_type="asearch", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_execute_search_without_proxy_auth_context_stays_sdk_only(monkeypatch): + """SDK callers have no key to attribute the search to, so no proxy metadata is invented.""" + import litellm + from litellm.proxy import proxy_server + + logger = WebSearchInterceptionLogger( + enabled_providers=["bedrock"], + search_tool_name="perplexity-sonar-pro", + ) + router = MagicMock() + router.search_tools = [ + { + "search_tool_name": "perplexity-sonar-pro", + "litellm_params": {"search_provider": "perplexity", "api_key": "fake-key"}, + } + ] + mock_asearch = AsyncMock(return_value=SearchResponse(object="search", results=[])) + + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "asearch", mock_asearch) + + await logger._execute_search("what is litellm", kwargs={"litellm_params": {}}) + + assert "litellm_metadata" not in mock_asearch.await_args.kwargs @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py index 544abab8dcf..c859f9b2f55 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py @@ -134,6 +134,30 @@ class TestBuildWebSearchToolResultBlock: assert first["title"] == "LiteLLM Docs" assert first["page_age"] == "2025-01-15" assert first["encrypted_content"] == "" + assert first["snippet"] == "Unified interface for LLMs." + + def test_snippet_carried_for_every_result(self): + # The snippet is the only field carrying page text. Losing it leaves the + # client and the model with nothing to answer from, forcing a fetch per + # result. + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=_make_search_response(), + ) + assert [r["snippet"] for r in block["content"]] == [ + "Unified interface for LLMs.", + "Pay-per-use pricing model.", + ] + + def test_missing_snippet_degrades_to_empty_string(self): + response = SearchResponse( + results=[SearchResult(title="T", url="https://x/", snippet="")] + ) + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=response, + ) + assert block["content"][0]["snippet"] == "" def test_handles_none_search_response(self): block = WebSearchTransformation.build_web_search_tool_result_block( @@ -223,11 +247,15 @@ class TestBuildPlanAttachesBlocks: ) blocks = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) - assert isinstance(blocks, list) - assert len(blocks) == 1 - assert blocks[0]["type"] == "web_search_tool_result" - assert blocks[0]["tool_use_id"] == "toolu_one" - assert blocks[0]["content"][0]["url"] == "https://docs.litellm.ai/" + assert isinstance(blocks, tuple) + assert [b["type"] for b in blocks] == [ + "server_tool_use", + "web_search_tool_result", + ] + assert blocks[0]["id"].startswith("srvtoolu_") + assert blocks[0]["input"] == {"query": "what is litellm"} + assert blocks[1]["tool_use_id"] == blocks[0]["id"] + assert blocks[1]["content"][0]["url"] == "https://docs.litellm.ai/" @pytest.mark.asyncio async def test_metadata_does_not_carry_blocks_when_flag_absent(self): @@ -479,6 +507,9 @@ class TestLegacyPathMatchesNewPath: kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, ) - assert out["content"][0]["type"] == "web_search_tool_result" - assert out["content"][0]["tool_use_id"] == "toolu_legacy" - assert out["content"][1]["type"] == "text" + assert [b["type"] for b in out["content"]] == [ + "server_tool_use", + "web_search_tool_result", + "text", + ] + assert out["content"][1]["tool_use_id"] == out["content"][0]["id"] diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 3728cc80323..b3956823dc1 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -723,6 +723,49 @@ class TestUnpackLegacyDefs: assert out["properties"]["r0"]["properties"]["p0"] == {"type": "string"} +class TestTextCompletionPromptToMessages: + """`/v1/completions` prompt wrapping, shared by the real-time and batch paths.""" + + def test_string_prompt_becomes_single_user_message(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + assert text_completion_prompt_to_messages("summarize this") == ( + {"role": "user", "content": "summarize this"}, + ) + + def test_list_of_strings_becomes_one_message_each(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + assert text_completion_prompt_to_messages(["first", "second"]) == ( + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ) + + @pytest.mark.parametrize( + "prompt", + [ + [1, 2, 3], + [[1, 2], [3, 4]], + ["ok", 7], + [], + "", + None, + {"role": "user"}, + ], + ) + def test_unsupported_prompt_shapes_raise(self, prompt): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + text_completion_prompt_to_messages, + ) + + with pytest.raises(ValueError, match="non-empty string or a non-empty list of strings"): + text_completion_prompt_to_messages(prompt) + + class TestCustomToolFormatShapeConversion: def test_flat_grammar_to_chat_shape(self): from litellm.litellm_core_utils.prompt_templates.common_utils import ( diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index 55db31efd2c..fb4cb494bee 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -162,3 +162,56 @@ class TestGetLitellmParamsDataResidency: api_base="https://eu.api.openai.com/v1", ) assert result["data_residency"] is None + + +class TestMetadataFallsBackToLitellmMetadata: + def test_metadata_falls_back_to_litellm_metadata_when_absent(self): + result = get_litellm_params(litellm_metadata={"trace_id": "trace-1"}) + assert result["metadata"] == {"trace_id": "trace-1"} + assert result["litellm_metadata"] == {"trace_id": "trace-1"} + + def test_empty_metadata_falls_back_to_litellm_metadata(self): + result = get_litellm_params(metadata={}, litellm_metadata={"trace_id": "trace-1"}) + assert result["metadata"] == {"trace_id": "trace-1"} + + def test_metadata_wins_when_both_present(self): + result = get_litellm_params( + metadata={"trace_id": "from-metadata"}, + litellm_metadata={"trace_id": "from-litellm-metadata"}, + ) + assert result["metadata"] == {"trace_id": "from-metadata"} + + @pytest.mark.parametrize("bad_value", ["not-json-a-string", 12345, ["a"], True]) + def test_non_dict_litellm_metadata_is_ignored(self, bad_value): + result = get_litellm_params(litellm_metadata=bad_value) + assert result["metadata"] is None + + def test_metadata_stays_none_without_litellm_metadata(self): + result = get_litellm_params(api_key="test-key") + assert result["metadata"] is None + + def test_session_and_trace_id_derived_from_litellm_metadata(self): + result = get_litellm_params( + litellm_metadata={"trace_id": "trace-1", "session_id": "session-1"}, + ) + assert result["litellm_session_id"] == "session-1" + assert result["litellm_trace_id"] == "trace-1" + + def test_explicit_session_and_trace_id_are_not_overridden(self): + result = get_litellm_params( + litellm_session_id="explicit-session", + litellm_trace_id="explicit-trace", + litellm_metadata={"trace_id": "trace-1", "session_id": "session-1"}, + ) + assert result["litellm_session_id"] == "explicit-session" + assert result["litellm_trace_id"] == "explicit-trace" + + def test_litellm_metadata_fallback_is_copied_not_aliased(self): + litellm_metadata = {"trace_id": "trace-1"} + + result = get_litellm_params(litellm_metadata=litellm_metadata) + + assert result["metadata"] == litellm_metadata + assert result["metadata"] is not litellm_metadata + result["metadata"].pop("trace_id") + assert litellm_metadata == {"trace_id": "trace-1"} 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 a09c45cb141..23e0975cd08 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -526,6 +526,26 @@ class TestUpdateFromKwargs: ) assert logging_obj.litellm_params["litellm_call_id"] == "call-empty" + @pytest.mark.parametrize("caller_metadata", [None, "not-a-dict", 42]) + def test_non_dict_caller_metadata_does_not_break_the_merge(self, logging_obj, caller_metadata): + logging_obj.update_from_kwargs( + kwargs={"metadata": caller_metadata, "litellm_metadata": {"user_api_key_hash": "hashed"}}, + litellm_params={"metadata": {"user_api_key_hash": "hashed", "litellm_api_version": "1.0"}}, + ) + + assert logging_obj.litellm_params["metadata"]["user_api_key_hash"] == "hashed" + + def test_does_not_mutate_caller_metadata_dict(self, logging_obj): + caller_metadata: dict = {} + + logging_obj.update_from_kwargs( + kwargs={"metadata": caller_metadata, "litellm_metadata": {"user_api_key_hash": "hashed"}}, + litellm_params={"metadata": {"user_api_key_hash": "hashed", "litellm_api_version": "1.0"}}, + ) + + assert caller_metadata == {} + assert logging_obj.litellm_params["metadata"]["user_api_key_hash"] == "hashed" + def test_logging_prevent_double_logging(logging_obj): """ diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py new file mode 100644 index 00000000000..3e85872f1e5 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_empty_choices.py @@ -0,0 +1,87 @@ +""" +Regression tests for OpenAI-compatible chunks with an empty ``choices`` list. + +``choices: []`` is valid OpenAI-compatible streaming: vLLM (and OpenAI itself, +when ``stream_options.include_usage`` is set) emits a final usage chunk with no +choices, and some gateways emit metadata-only chunks mid-stream. The adapter +used to index ``chunk.choices[0]`` unconditionally, so such a chunk raised +``IndexError: list index out of range`` and killed the ``/v1/messages`` stream. +""" + +import asyncio +import json +from typing import Any, AsyncIterator, Dict, List, Optional + +from litellm.llms.anthropic.experimental_pass_through.adapters.streaming_iterator import ( + AnthropicStreamWrapper, +) +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + +def _text_chunk(text: str) -> ModelResponseStream: + return ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content=text), finish_reason=None)] + ) + + +def _finish_chunk() -> ModelResponseStream: + return ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]) + + +def _empty_choices_chunk(usage: Optional[Usage] = None) -> ModelResponseStream: + return ModelResponseStream(choices=[], usage=usage) + + +def _collect_async(wrapper: AnthropicStreamWrapper) -> str: + async def _run() -> str: + return "".join( + [raw.decode() if isinstance(raw, bytes) else raw async for raw in wrapper.async_anthropic_sse_wrapper()] + ) + + return asyncio.run(_run()) + + +def _message_delta(sse: str) -> Dict[str, Any]: + return next( + json.loads(line[len("data: ") :]) + for block in sse.split("\n\n") + for line in block.splitlines() + if line.startswith("data: ") and '"message_delta"' in line + ) + + +def test_leading_metadata_chunk_without_choices_does_not_kill_stream(): + """A metadata-only chunk before any content must be skipped, not indexed.""" + chunks: List[ModelResponseStream] = [ + _empty_choices_chunk(), + _text_chunk("Hello"), + _text_chunk(" there"), + _finish_chunk(), + ] + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="mock-model") + events = list(wrapper) + + text = "".join( + event["delta"]["text"] for event in events if event.get("type") == "content_block_delta" + ) + assert text == "Hello there" + assert events[-1]["type"] == "message_stop" + + +def test_final_usage_chunk_without_choices_is_merged_into_message_delta(): + """The vLLM/OpenAI final usage chunk carries no choices; its usage must + still land on the Anthropic ``message_delta``.""" + usage = Usage(prompt_tokens=10, completion_tokens=3, total_tokens=13) + + async def _aiter() -> "AsyncIterator[ModelResponseStream]": + for chunk in [_text_chunk("Hi"), _finish_chunk(), _empty_choices_chunk(usage)]: + yield chunk + + sse = _collect_async(AnthropicStreamWrapper(completion_stream=_aiter(), model="mock-model")) + + message_delta = _message_delta(sse) + assert message_delta["delta"]["stop_reason"] == "end_turn" + assert message_delta["usage"]["input_tokens"] == 10 + assert message_delta["usage"]["output_tokens"] == 3 + assert "Hi" in sse + assert "message_stop" in sse diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index df3db3d2c57..f11324ca376 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -732,6 +732,61 @@ def test_handler_skips_strip_when_presanitized(): assert result is not None +def test_handler_flattens_replayed_unencrypted_web_search_results(): + """Synthesized search blocks replayed as history must reach the provider as text.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + captured = {} + + def fake_base_handler(*args, **kwargs): + captured.update(kwargs) + return "stub" + + with patch.object( + handler.base_llm_http_handler, + "anthropic_messages_handler", + side_effect=fake_base_handler, + ): + handler.anthropic_messages_handler( + max_tokens=10, + messages=[ + {"role": "user", "content": "latest litellm version?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "latest litellm version"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://github.com/BerriAI/litellm/releases", + "title": "Releases", + "page_age": None, + "encrypted_content": "", + "snippet": "Latest release v1.95.0", + } + ], + }, + ], + }, + {"role": "user", "content": "which version?"}, + ], + model="anthropic/claude-3-5-sonnet-20241022", + custom_llm_provider="anthropic", + ) + + replayed = captured["messages"][1]["content"] + assert [b["type"] for b in replayed] == ["text"] + assert "Snippet: Latest release v1.95.0" in replayed[0]["text"] + + def test_presanitized_flag_not_leaked_to_provider_params(): """The private sentinel must be popped, never forwarded as a request param.""" from litellm.llms.anthropic.experimental_pass_through.messages import handler diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 6ab0f2c08ab..9df72108332 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -10,8 +10,10 @@ Verifies that: - ANTHROPIC_API_KEY / ANTHROPIC_API_BASE take precedence over their aliases. """ +import json import os import sys +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -1457,6 +1459,190 @@ class TestAnthropicThinkingSignatureSelfHeal: out = strip_empty_text_blocks_from_anthropic_messages(msgs) assert [b["type"] for b in out[0]["content"]] == ["tool_result"] + def test_flatten_unencrypted_web_search_results_keeps_snippet_evidence(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + {"role": "user", "content": "latest litellm version?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "latest litellm version"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://github.com/BerriAI/litellm/releases", + "title": "Releases", + "page_age": None, + "encrypted_content": "", + "snippet": "Latest release v1.95.0", + } + ], + }, + {"type": "text", "text": "v1.95.0"}, + ], + }, + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + assert [b["type"] for b in out[1]["content"]] == ["text", "text"] + flattened = out[1]["content"][0]["text"] + assert "Web search results for 'latest litellm version':" in flattened + assert "URL: https://github.com/BerriAI/litellm/releases" in flattened + assert "Snippet: Latest release v1.95.0" in flattened + assert msgs[1]["content"][0]["type"] == "server_tool_use" + + @pytest.mark.parametrize("results", [[], None], ids=["empty_list", "search_raised"]) + def test_flatten_unencrypted_web_search_results_flattens_a_resultless_search(self, results): + """A search that found nothing, or that raised, still has to be flattened. + + Both cases reach the client as ``content: []``, and leaving that block in + place ships an unsupported tag to Bedrock on the next turn just as surely + as a populated one does. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=None if results is None else SimpleNamespace(results=results), + ) + assert block["content"] == [], "fixture drifted from what the interceptor emits" + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "who won"}, + }, + block, + {"type": "text", "text": "I could not find that."}, + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert [b["type"] for b in out[0]["content"]] == ["text", "text"] + assert out[0]["content"][0]["text"] == ("Web search results for 'who won':\n\nNo results were returned.") + + @pytest.mark.parametrize("results", [[SimpleNamespace(title="Rome", url="u", snippet="s", date=None)], []]) + def test_flatten_unencrypted_web_search_results_is_idempotent(self, results): + """Flattening twice must equal flattening once. + + The agentic loop re-enters the same entry point for its follow-up call and + hands it the original history, so this runs again on already-flattened + messages once per iteration. A pass that appended instead of replacing + would duplicate the evidence on every loop. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "when"}}, + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=SimpleNamespace(results=results), + ), + {"type": "text", "text": "753 BC."}, + ], + } + ] + + once = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + twice = flatten_unencrypted_web_search_results_in_anthropic_messages(once) + + assert [b["type"] for b in once[0]["content"]] == ["text", "text"] + assert json.dumps(twice) == json.dumps(once) + + def test_flatten_unencrypted_web_search_results_preserves_real_anthropic_blocks(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "q"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://example.com", + "title": "Example", + "page_age": None, + "encrypted_content": "EqgfCioIARgBIiQ4", + } + ], + }, + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + + def test_flatten_unencrypted_web_search_results_leaves_error_blocks_alone(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "web_search_tool_result_error", + "error_code": "max_uses_exceeded", + }, + } + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + def test_sanitize_tool_use_ids_in_anthropic_messages(self): from litellm.llms.anthropic.common_utils import ( sanitize_tool_use_ids_in_anthropic_messages, diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 3681daffe5e..01420eb10df 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -172,7 +172,7 @@ def test_create_request_omits_kms_key_when_absent(config): "generate_unique_job_name", return_value="litellm-batch-1", ), patch.object(config.common_utils, "sign_aws_request") as mock_sign, patch( - "litellm.llms.bedrock.batches.transformation.get_secret_str", + "litellm.llms.bedrock.common_utils.get_secret_str", return_value=None, ): mock_sign.return_value = ({}, b"{}") diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 2ec0c00db26..270add48e0e 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -446,7 +446,7 @@ class TestBedrockFilesTransformation: captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -502,7 +502,7 @@ class TestBedrockFilesTransformation: captured_optional_params: dict = {} - def fake_sign(content, api_base, optional_params): + def fake_sign(content, api_base, optional_params, s3_encryption_key_id=None): captured_optional_params.update(optional_params) return {"Authorization": "fake"}, content @@ -518,6 +518,74 @@ class TestBedrockFilesTransformation: captured_optional_params.get("aws_region_name") == "us-gov-west-1" ), "s3_region_name must override aws_region_name for SigV4 signing" + def _signed_upload_request(self, litellm_params: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + config = BedrockFilesConfig() + jsonl_content = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + }, + } + ).encode() + + request = config.transform_create_file_request( + model="amazon.nova-pro-v1:0", + create_file_data={ + "file": ("batch.jsonl", jsonl_content, "application/jsonl"), + "purpose": "batch", + }, + optional_params={ + "aws_access_key_id": "test-key-id", + "aws_secret_access_key": "test-secret", + "aws_region_name": "us-west-2", + }, + litellm_params={"s3_bucket_name": "litellm-batch-bucket", **litellm_params}, + ) + assert isinstance(request, dict) + return request + + def test_upload_signs_sse_kms_headers_when_key_configured(self, monkeypatch): + """ + Buckets whose policy requires SSE-KMS reject the batch input-file PutObject + unless the upload carries the aws:kms encryption headers; they must also be + covered by SigV4 SignedHeaders or S3 answers SignatureDoesNotMatch. + """ + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + kms_key = "arn:aws:kms:us-west-2:1234:key/abcd" + + request = self._signed_upload_request({"s3_encryption_key_id": kms_key}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption"] == "aws:kms" + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == kms_key + signed_headers = headers["authorization"].split("SignedHeaders=")[1].split(",")[0] + assert "x-amz-server-side-encryption" in signed_headers + assert "x-amz-server-side-encryption-aws-kms-key-id" in signed_headers + + def test_upload_reads_sse_kms_key_from_env(self, monkeypatch): + monkeypatch.setenv("AWS_S3_ENCRYPTION_KEY_ID", "env-kms-key") + + request = self._signed_upload_request({}) + + headers = {key.lower(): value for key, value in request["headers"].items()} + assert headers["x-amz-server-side-encryption-aws-kms-key-id"] == "env-kms-key" + + def test_upload_omits_sse_headers_when_no_key_configured(self, monkeypatch): + monkeypatch.delenv("AWS_S3_ENCRYPTION_KEY_ID", raising=False) + + request = self._signed_upload_request({}) + + headers = {key.lower() for key in request["headers"]} + assert "x-amz-server-side-encryption" not in headers + assert "x-amz-server-side-encryption-aws-kms-key-id" not in headers + def test_openai_passthrough_still_works(self): """ Regression test: ensure OpenAI-compatible models (e.g. gpt-oss) @@ -1076,21 +1144,64 @@ class TestBedrockFilesEmbeddingTransformation: is None ) - def test_is_embedding_record_helper(self): - """Helper detects embeddings via `url` first, then by body shape.""" + def test_classify_batch_record_helper(self): + """Helper classifies by `url` first, then by body shape.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind - assert BedrockFilesConfig._is_embedding_record( - {"url": "/v1/embeddings", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/embeddings", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.EMBEDDING ) # body-only fallback - assert BedrockFilesConfig._is_embedding_record({"body": {"input": "x"}}) - # chat shape - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/chat/completions", "body": {"messages": []}} + assert ( + BedrockFilesConfig._classify_batch_record({"body": {"input": "x"}}) + is BedrockBatchRecordKind.EMBEDDING ) - # ambiguous body without `input` is treated as not-embedding - assert not BedrockFilesConfig._is_embedding_record({"body": {}}) + # chat shape + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/chat/completions", "body": {"messages": []}} + ) + is BedrockBatchRecordKind.CHAT + ) + # ambiguous body without any recognized key is treated as chat + assert ( + BedrockFilesConfig._classify_batch_record({"body": {}}) + is BedrockBatchRecordKind.CHAT + ) + + @pytest.mark.parametrize("body", ["not a mapping", ["messages"], 7, None], ids=["str", "list", "int", "missing"]) + def test_classify_batch_record_falls_back_to_chat_for_non_mapping_body(self, body): + """A malformed body must not crash the whole upload during classification. + + Chat is the only kind whose transformer tolerates an unexpected shape and + raises a readable error; routing a non-mapping body anywhere else would + blow up on attribute access before the caller sees which record is bad. + """ + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind + + record = {"body": body} if body is not None else {} + assert BedrockFilesConfig._classify_batch_record(record) is BedrockBatchRecordKind.CHAT + + def test_embedding_kind_is_rejected_by_the_chat_normalizer(self): + """Embeddings have no chat equivalent, so the normalizer refuses them outright. + + The caller routes embeddings to the Titan transformer before ever getting + here; this guard is what keeps a future caller from quietly shipping an + embedding body through the chat path. + """ + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind + + with pytest.raises(ValueError, match="do not have a chat-completion equivalent"): + BedrockFilesConfig._transform_batch_body_to_chat_body( + {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": "hi"}, + BedrockBatchRecordKind.EMBEDDING, + ) def test_explicit_chat_url_with_input_body_short_circuits_to_chat(self): """Explicit url=/v1/chat/completions wins even if body looks like embedding. @@ -1101,15 +1212,20 @@ class TestBedrockFilesEmbeddingTransformation: """ from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind + # Direct helper assertion - assert not BedrockFilesConfig._is_embedding_record( - { - "url": "/v1/chat/completions", - "body": { - "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - "input": "this would mis-route under the old precedence", - }, - } + assert ( + BedrockFilesConfig._classify_batch_record( + { + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "input": "this would mis-route under the old precedence", + }, + } + ) + is BedrockBatchRecordKind.CHAT ) # End-to-end: a record like this routes through the chat path. We @@ -1168,20 +1284,347 @@ class TestBedrockFilesEmbeddingTransformation: with pytest.raises(ValueError, match="must be a string"): BedrockFilesConfig._coerce_embedding_input_to_string({"unsupported": True}) - def test_other_non_embedding_urls_route_to_chat(self): - """Any non-/v1/embeddings url short-circuits to chat path.""" + def test_other_non_embedding_urls_do_not_route_to_embeddings(self): + """An `input` body only means "embedding" when the url says so.""" from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + from litellm.types.llms.bedrock import BedrockBatchRecordKind # /v1/completions (legacy completions endpoint) - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/completions", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/completions", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.TEXT_COMPLETION + ) + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/responses", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.RESPONSES ) # Arbitrary unknown url - caller's explicit signal still wins - assert not BedrockFilesConfig._is_embedding_record( - {"url": "/v1/responses", "body": {"input": "x"}} + assert ( + BedrockFilesConfig._classify_batch_record( + {"url": "/v1/moderations", "body": {"input": "x"}} + ) + is BedrockBatchRecordKind.CHAT ) +class TestBedrockBatchNonChatEndpointRecords: + """`/v1/completions` and `/v1/responses` JSONL records (issue #35639). + + Bedrock batch `modelInput` is always the model's InvokeModel/Converse body, + so a record shaped for another OpenAI endpoint has to be normalized to chat + completions first. Before this normalization every record below either + raised `BadRequestError` at `POST /v1/files` (Anthropic, Nova) or silently + shipped an empty `messages` list to AWS (passthrough providers). + """ + + ANTHROPIC_MODEL = "bedrock/us.anthropic.claude-sonnet-4-6" + NOVA_MODEL = "bedrock/us.amazon.nova-pro-v1:0" + PASSTHROUGH_MODEL = "bedrock/openai.gpt-oss-120b-1:0" + + def _transform(self, record: dict) -> dict: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content([record]) + assert len(result) == 1 + assert result[0]["recordId"] == record["custom_id"] + return result[0]["modelInput"] + + def test_anthropic_text_completion_record_wraps_prompt(self): + model_input = self._transform( + { + "custom_id": "1", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "prompt": "Summarize the following call transcript", + "max_tokens": 64, + }, + } + ) + + assert model_input == { + "messages": [ + { + "role": "user", + "content": [{"type": "text", "text": "Summarize the following call transcript"}], + } + ], + "max_tokens": 64, + "anthropic_version": "bedrock-2023-05-31", + } + + def test_anthropic_text_completion_record_keeps_every_prompt_in_a_list(self): + model_input = self._transform( + { + "custom_id": "2", + "method": "POST", + "url": "/v1/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "prompt": ["first prompt", "second prompt"], + "max_tokens": 8, + }, + } + ) + + # Consecutive user messages are merged by the Anthropic transform, the + # same way they are on the real-time path. + assert model_input["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "first prompt"}, + {"type": "text", "text": "second prompt"}, + ], + } + ] + assert "prompt" not in model_input + + def test_anthropic_responses_record_wraps_string_input(self): + model_input = self._transform( + { + "custom_id": "3", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "input": "hi", + "max_output_tokens": 16, + }, + } + ) + + assert model_input == { + "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + "max_tokens": 16, + "anthropic_version": "bedrock-2023-05-31", + } + assert "tools" not in model_input, "an empty tools array must not be shipped to Bedrock" + + def test_anthropic_responses_record_maps_instructions_and_input_items(self): + """The Responses-specific params go through the same bridge as real time.""" + model_input = self._transform( + { + "custom_id": "4", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.ANTHROPIC_MODEL, + "instructions": "be terse", + "input": [ + {"role": "user", "content": "what is 2+2?"}, + {"role": "assistant", "content": "4"}, + {"role": "user", "content": "and 3+3?"}, + ], + "max_output_tokens": 32, + "temperature": 0.2, + }, + } + ) + + assert model_input["system"] == [{"type": "text", "text": "be terse"}] + assert model_input["max_tokens"] == 32 + assert model_input["temperature"] == 0.2 + assert [message["role"] for message in model_input["messages"]] == [ + "user", + "assistant", + "user", + ] + assert model_input["messages"][-1]["content"] == [{"type": "text", "text": "and 3+3?"}] + assert "input" not in model_input + assert "max_output_tokens" not in model_input + + def test_responses_record_keeps_metadata(self): + """`metadata` reaches the bridge, which reads it as its own kwarg.""" + model_input = self._transform( + { + "custom_id": "4b", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": self.PASSTHROUGH_MODEL, + "input": "hi", + "metadata": {"tenant": "acct-1"}, + }, + } + ) + + assert model_input["metadata"] == {"tenant": "acct-1"} + + @pytest.mark.parametrize("model_attr", ["ANTHROPIC_MODEL", "NOVA_MODEL"], ids=["anthropic", "nova"]) + def test_modelled_providers_do_not_smuggle_metadata_into_the_bedrock_body(self, model_attr): + """Providers with a real InvokeModel schema leave `metadata` out of `modelInput`. + + Batch `modelInput` has to match the model's own InvokeModel body, and + neither the Anthropic messages body nor the Nova body has a field for + arbitrary caller labels. Nova in particular answers `400 Malformed input + request` for any key it does not recognize, so translating `metadata` + into the Converse-level `requestMetadata` would fail the record rather + than preserve the labels. The passthrough providers keep it because + their body is the OpenAI request itself. + """ + model_input = self._transform( + { + "custom_id": "4c", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": getattr(self, model_attr), + "input": "hi", + "max_output_tokens": 8, + "metadata": {"tenant": "acct-1"}, + }, + } + ) + + assert "metadata" not in model_input + assert "requestMetadata" not in model_input + + @pytest.mark.parametrize( + "body", + [ + {"prompt": "hi"}, + {"input": "hi"}, + ], + ids=["prompt", "input"], + ) + def test_nova_converse_record_wraps_prompt_and_input(self, body): + url = "/v1/completions" if "prompt" in body else "/v1/responses" + model_input = self._transform( + { + "custom_id": "5", + "method": "POST", + "url": url, + "body": {"model": self.NOVA_MODEL, **body}, + } + ) + + assert model_input["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + + @pytest.mark.parametrize( + "body", + [ + {"prompt": "hi"}, + {"input": "hi"}, + ], + ids=["prompt", "input"], + ) + def test_passthrough_provider_record_no_longer_emits_empty_messages(self, body): + """The passthrough branch used to emit `{"messages": [], "prompt": ...}`. + + That shape is accepted by `POST /v1/files`, so the whole batch job was + submitted to AWS and only failed there. + """ + url = "/v1/completions" if "prompt" in body else "/v1/responses" + model_input = self._transform( + { + "custom_id": "6", + "method": "POST", + "url": url, + "body": {"model": self.PASSTHROUGH_MODEL, **body}, + } + ) + + # Asserted on the serialized form, since the passthrough branch hands + # `messages` straight to S3 without a per-provider transform. + assert json.loads(json.dumps(model_input)) == {"messages": [{"role": "user", "content": "hi"}]} + + def test_mixed_endpoints_in_one_file_keep_their_own_shapes(self): + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + result = BedrockFilesConfig()._transform_openai_jsonl_content_to_bedrock_jsonl_content( + [ + { + "custom_id": "chat", + "url": "/v1/chat/completions", + "body": { + "model": self.ANTHROPIC_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 4, + }, + }, + { + "custom_id": "text", + "url": "/v1/completions", + "body": {"model": self.ANTHROPIC_MODEL, "prompt": "hi", "max_tokens": 4}, + }, + { + "custom_id": "responses", + "url": "/v1/responses", + "body": {"model": self.ANTHROPIC_MODEL, "input": "hi", "max_output_tokens": 4}, + }, + { + "custom_id": "embedding", + "url": "/v1/embeddings", + "body": {"model": "bedrock/amazon.titan-embed-text-v2:0", "input": "hi"}, + }, + ] + ) + + assert [record["recordId"] for record in result] == [ + "chat", + "text", + "responses", + "embedding", + ] + for record in result[:3]: + assert record["modelInput"]["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]} + ] + assert result[3]["modelInput"] == {"inputText": "hi"} + + @pytest.mark.parametrize( + ("url", "expected_message"), + [ + ("/v1/completions", "missing required `prompt` field"), + ("/v1/responses", "missing required `input` field"), + ], + ) + def test_missing_required_field_raises_actionable_error(self, url, expected_message): + with pytest.raises(ValueError, match=expected_message): + self._transform( + { + "custom_id": "7", + "method": "POST", + "url": url, + "body": {"model": self.ANTHROPIC_MODEL, "max_tokens": 4}, + } + ) + + def test_prompt_body_without_url_is_still_wrapped(self): + """A record can omit `url`; the body shape then decides.""" + model_input = self._transform( + { + "custom_id": "8", + "body": {"model": self.ANTHROPIC_MODEL, "prompt": "hi", "max_tokens": 4}, + } + ) + + assert model_input["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + + def test_messages_win_over_prompt_when_url_is_absent(self): + model_input = self._transform( + { + "custom_id": "9", + "body": { + "model": self.ANTHROPIC_MODEL, + "messages": [{"role": "user", "content": "from messages"}], + "prompt": "from prompt", + "max_tokens": 4, + }, + } + ) + + assert model_input["messages"] == [ + {"role": "user", "content": [{"type": "text", "text": "from messages"}]} + ] + + class TestBedrockFileContentTransformation: """SigV4-signed S3 GetObject retrieval of Bedrock batch output files.""" 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 3b8b4af78d9..76bb11cc26d 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 @@ -4,6 +4,7 @@ import json import os import sys from datetime import datetime +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -2515,3 +2516,67 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( assert thinking.get("type") == "enabled" assert isinstance(thinking.get("budget_tokens"), int) assert "output_config" not in flipped + + +@pytest.mark.parametrize( + "search_results, expected_evidence", + [ + pytest.param( + [SimpleNamespace(title="Rome", url="https://ex.com/rome", snippet="Founded 753 BC.", date=None)], + "Snippet: Founded 753 BC.", + id="search_returned_results", + ), + pytest.param([], "No results were returned.", id="search_returned_nothing"), + ], +) +def test_replayed_intercepted_search_turn_leaves_no_unsupported_block_for_bedrock(search_results, expected_evidence): + """A native client replaying an intercepted search turn must not 400 on Bedrock. + + ``websearch_interception`` hands Claude Desktop an Anthropic-native + ``server_tool_use`` + ``web_search_tool_result`` pair, and Anthropic's protocol + obliges the client to replay that assistant turn verbatim on every later turn. + Bedrock's Anthropic schema defines neither tag, so both have to be gone from the + outbound body by the time it is signed, with the search evidence carried forward + as text instead. Built from the real builder rather than a hand-written fixture + so the two cannot drift apart. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + from litellm.types.router import GenericLiteLLMParams + + replayed_turn = [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "when was Rome founded"}, + }, + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=SimpleNamespace(results=search_results), + ), + {"type": "text", "text": "Rome was founded in 753 BC."}, + ] + messages = [ + {"role": "user", "content": [{"type": "text", "text": "When was Rome founded?"}]}, + {"role": "assistant", "content": replayed_turn}, + {"role": "user", "content": [{"type": "text", "text": "Repeat the year."}]}, + ] + + body = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=flatten_unencrypted_web_search_results_in_anthropic_messages(messages), + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + serialized = json.dumps(body) + assert "web_search_tool_result" not in serialized + assert "server_tool_use" not in serialized + assert expected_evidence in serialized + assert "Rome was founded in 753 BC." in serialized diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py index f279acfd60c..82010e82cea 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py @@ -13,7 +13,7 @@ def test_create_aiohttp_transport_sets_enable_cleanup_closed_when_needed(monkeyp ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport( shared_session=None ) @@ -36,7 +36,7 @@ def test_create_aiohttp_transport_omits_enable_cleanup_closed_when_not_needed( ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport( shared_session=None ) diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py index 5a37e681c4a..0065bf8f4ef 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_so_keepalive.py @@ -34,7 +34,7 @@ def test_socket_factory_omitted_when_disabled(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 @@ -55,7 +55,7 @@ def test_socket_factory_attached_when_enabled(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 @@ -77,7 +77,7 @@ def test_socket_factory_skipped_on_old_aiohttp(monkeypatch): ) as mock_tcp_connector: with patch.object( http_handler_module, "ClientSession", return_value=session_mock - ): + ), patch.object(http_handler_module, "DummyCookieJar", return_value=MagicMock(name="cookie_jar")): _invoke_connector_factory(http_handler_module) assert mock_tcp_connector.call_count >= 1 diff --git a/tests/test_litellm/llms/custom_httpx/test_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_http_handler.py index 86d097c6123..fa1c7308c6f 100644 --- a/tests/test_litellm/llms/custom_httpx/test_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_http_handler.py @@ -1171,3 +1171,77 @@ async def test_client_handed_out_by_async_cache_survives_eviction_and_collection assert not consumer_client.is_closed await consumer_client.aclose() + + +_SET_COOKIE = "SESSION=upstream-a-secret; Path=/" + + +def _cookie_recorder(): + """A transport that hands out a Set-Cookie once, and records what comes back.""" + seen = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen.append(request.headers.get("cookie")) + if request.url.path == "/set": + return httpx.Response(200, headers={"set-cookie": _SET_COOKIE}) + return httpx.Response(200) + + return handler, seen + + +@pytest.mark.asyncio +async def test_async_client_never_replays_one_upstreams_cookie_to_another(): + """LiteLLM's async clients are pooled and shared by every caller, so a cookie one + upstream sets would be attached to every later request on a matching domain, reaching + a different tenant's upstream. The client must persist no response cookie.""" + handler, seen = _cookie_recorder() + http_handler = AsyncHTTPHandler() + client = http_handler.client + client._transport = httpx.MockTransport(handler) + + await client.get("https://upstream-a.example.com/set") + await client.get("https://upstream-b.example.com/rpc") + await client.aclose() + + assert dict(client.cookies) == {}, "the shared client stored an upstream's cookie" + assert seen == [None, None] + + +def test_sync_client_never_replays_one_upstreams_cookie_to_another(): + """Same invariant on the sync client, which is pooled the same way.""" + handler, seen = _cookie_recorder() + http_handler = HTTPHandler() + client = http_handler.client + client._transport = httpx.MockTransport(handler) + + client.get("https://upstream-a.example.com/set") + client.get("https://upstream-b.example.com/rpc") + client.close() + + assert dict(client.cookies) == {} + assert seen == [None, None] + + +@pytest.mark.asyncio +async def test_aiohttp_session_never_replays_one_upstreams_cookie_to_another(): + """The httpx jar is not the only one. AiohttpTransport is litellm's default transport + and the aiohttp ClientSession keeps its own cookie jar, which httpx-level assertions + cannot see, so blocking only the httpx jar leaves the leak intact on the real path. + + aiohttp's default jar refuses cookies for IP hosts, so this drives a hostname. An + IP-addressed check passes whether or not the session jar is blocked.""" + from aiohttp import DummyCookieJar + from yarl import URL + + http_handler = AsyncHTTPHandler(timeout=61.0) + transport = http_handler.client._transport + assert isinstance(transport, LiteLLMAiohttpTransport), "aiohttp is no longer the default transport" + + session = transport.client() if callable(transport.client) else transport.client + jar = session.cookie_jar + assert isinstance(jar, DummyCookieJar) + + jar.update_cookies({"SESSION": "upstream-a-secret"}, URL("https://upstream-a.example.com")) + assert len(jar) == 0 + assert dict(jar.filter_cookies(URL("https://upstream-a.example.com"))) == {} + await session.close() diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index cacea234777..9535bf17411 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -5,8 +5,10 @@ The handler is HTTP/auth glue around the (separately-tested) pure ``VertexAIBatchTransformation``. Each public method (create / retrieve / list / cancel) resolves a Vertex access token + URL, branches on ``_is_async`` (returning the coroutine in the async case, doing the sync HTTP call otherwise), -checks the HTTP status, and parses the JSON into ``LiteLLMBatch`` (or the OpenAI -list shape). +and parses the JSON into ``LiteLLMBatch`` (or the OpenAI list shape). POST-backed +calls rely on the client's ``raise_for_status`` (non-2xx surfaces as +``httpx.HTTPStatusError``); GET-backed calls return without raising, so the +handler checks their status codes itself. We mock only true I/O / auth seams: * ``_ensure_access_token`` - the Vertex credential seam. Returns a fixed @@ -20,7 +22,7 @@ We mock only true I/O / auth seams: what URL/headers/body, and that the response is parsed into the litellm type. Sibling seams are asserted NOT called where relevant. -The ``_is_async`` branch, status-code error paths, and the cancel +The ``_is_async`` branch, the error paths, and the cancel retrieve-after-cancel sequencing run for real. """ @@ -40,6 +42,7 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) +from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 HMOD = "litellm.llms.vertex_ai.batches.handler" @@ -178,13 +181,19 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): sync_client.post.assert_not_called() -def test_create_batch_sync_non_200_raises(): +def test_create_batch_sync_httpstatuserror_propagates(): + """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the + sync create path must surface that error, not swallow it.""" h = _make_handler() client = MagicMock() - client.post.return_value = _http_response(status_code=500) + request = httpx.Request("POST", "https://x/batchPredictionJobs") + err_response = httpx.Response(status_code=500, request=request, text="boom") + client.post.side_effect = httpx.HTTPStatusError( + "boom", request=request, response=err_response + ) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(httpx.HTTPStatusError): h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -197,27 +206,27 @@ def test_create_batch_sync_non_200_raises(): ) -def test_create_batch_async_non_200_raises(): +def test_create_batch_input_file_id_without_model_raises_400_before_post(): + """A gs:// uri with no publishers//models/ path is a 400, not a bare 500.""" h = _make_handler() - async_client = MagicMock() - async_client.post = AsyncMock(return_value=_http_response(status_code=403)) + client = MagicMock() - with ( - patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), - patch(f"{HMOD}.get_async_httpx_client", return_value=async_client), - ): - coro = h.create_batch( - _is_async=True, - create_batch_data=CREATE_DATA, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - with pytest.raises(Exception, match="Error: 403"): - _run(coro) + with patch(f"{HMOD}._get_httpx_client", return_value=client): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data={"input_file_id": "gs://bucket/batch-input.jsonl"}, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + assert "gs://bucket/batch-input.jsonl" in str(exc_info.value) + client.post.assert_not_called() # =========================================================================== # @@ -292,7 +301,7 @@ def test_retrieve_batch_sync_non_200_raises(): patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), patch(f"{HMOD}.safe_get", return_value=_http_response(status_code=404)), ): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.retrieve_batch( _is_async=False, batch_id=BATCH_ID, @@ -438,7 +447,7 @@ def test_list_batches_sync_non_200_raises(): client.get.return_value = _http_response(status_code=500) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): h.list_batches( _is_async=False, after=None, @@ -524,27 +533,6 @@ def test_cancel_batch_async_returns_coroutine_posts_then_retrieves(): assert post_kwargs["url"].endswith(":cancel") -def test_cancel_batch_sync_cancel_post_non_200_raises(): - h = _make_handler() - client = MagicMock() - client.post.return_value = _http_response(status_code=500) - - with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 500"): - h.cancel_batch( - _is_async=False, - batch_id=BATCH_ID, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - # cancel POST failed -> retrieve GET must never fire - client.get.assert_not_called() - - def test_cancel_batch_sync_retrieve_non_200_raises(): h = _make_handler() client = MagicMock() @@ -552,7 +540,7 @@ def test_cancel_batch_sync_retrieve_non_200_raises(): client.get.return_value = _http_response(status_code=404) with patch(f"{HMOD}._get_httpx_client", return_value=client): - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): h.cancel_batch( _is_async=False, batch_id=BATCH_ID, @@ -672,7 +660,7 @@ def test_async_retrieve_batch_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -726,7 +714,7 @@ def test_async_list_batches_non_200_raises(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 500"): + with pytest.raises(VertexAIError, match="Error: 500"): _run(coro) @@ -761,28 +749,6 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): _run(coro) async_client.get.assert_not_awaited() - # (a2) cancel POST returns a plain non-200 (no exception) -> raises - async_client_post500 = MagicMock() - async_client_post500.post = AsyncMock(return_value=_http_response(status_code=500)) - async_client_post500.get = AsyncMock() - with ( - patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), - patch(f"{HMOD}.get_async_httpx_client", return_value=async_client_post500), - ): - coro = h.cancel_batch( - _is_async=True, - batch_id=BATCH_ID, - api_base=None, - vertex_credentials=None, - vertex_project=PROJECT, - vertex_location=LOCATION, - timeout=600.0, - max_retries=None, - ) - with pytest.raises(Exception, match="Error: 500"): - _run(coro) - async_client_post500.get.assert_not_awaited() - # (b) retrieve-after-cancel returns non-200 async_client2 = MagicMock() async_client2.post = AsyncMock(return_value=_http_response(json_body={})) @@ -801,5 +767,5 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): timeout=600.0, max_retries=None, ) - with pytest.raises(Exception, match="Error: 404"): + with pytest.raises(VertexAIError, match="Error: 404"): _run(coro) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 1b37ade6b30..8352ec16389 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -25,6 +25,7 @@ from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, ) from litellm.llms.vertex_ai.common_utils import ( # noqa: E402 + VertexAIError, _convert_vertex_datetime_to_openai_datetime, ) from litellm.types.utils import LiteLLMBatch # noqa: E402 @@ -69,6 +70,24 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +@pytest.mark.parametrize( + "input_file_id", + [ + "gs://bucket/no-model-here.jsonl", + "gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", + "gs://bucket/publishers/google/models", + "gs://bucket/publishers/google/models//file-uuid", + ], +) +def test_transform_openai_request_unparseable_model_raises_400(input_file_id: str): + """An input_file_id with no parseable model path is a client error, not an IndexError -> 500.""" + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": input_file_id}) + + assert exc_info.value.status_code == 400 + assert input_file_id in str(exc_info.value) + + # =========================================================================== # # transform_vertex_ai_batch_response_to_openai_batch_response # =========================================================================== # @@ -299,9 +318,29 @@ def test_get_model_from_gcs_file_url_encoded(): assert T._get_model_from_gcs_file(encoded) == "publishers/google/models/gemini-1.5-flash-001" -def test_get_model_from_gcs_file_no_publishers_raises(): - with pytest.raises(IndexError): +def test_get_model_from_gcs_file_no_publishers_raises_400(): + with pytest.raises(VertexAIError) as exc_info: T._get_model_from_gcs_file("gs://bucket/no-model-here.jsonl") + assert exc_info.value.status_code == 400 + + +# =========================================================================== # +# is_unmanaged_gcs_batch_input_file_id +# =========================================================================== # + + +@pytest.mark.parametrize( + "input_file_id, expected", + [ + (INPUT_FILE, True), + (None, False), + ("file-abc123", False), + ("gs://bucket/no-model-here.jsonl", False), + ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + ], +) +def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): + assert T.is_unmanaged_gcs_batch_input_file_id(input_file_id) is expected # =========================================================================== # diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 6bdea9c2615..066d33dc187 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -2,8 +2,11 @@ Unit tests for AgentRequestHandler - Agent permission management for keys and teams. """ +import hashlib +import json import os import sys +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -11,8 +14,11 @@ import pytest sys.path.insert(0, os.path.abspath("../../../..")) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) @@ -22,11 +28,11 @@ class TestAgentRequestHandler: Test suite for AgentRequestHandler permission logic. """ - async def test_get_allowed_agents_intersection_logic(self): + async def test_resolve_agent_access_intersection_logic(self): """ Test key/team intersection: when both have restrictions, only common agents are allowed. When team has restrictions but key has none, key inherits from team. - When neither has restrictions, returns empty list (meaning allow all). + Only a caller with no grant anywhere is unrestricted. """ mock_user_auth = UserAPIKeyAuth( api_key="test-key", @@ -41,13 +47,17 @@ class TestAgentRequestHandler: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: - mock_key.return_value = ["agent1", "agent2", "agent3"] - mock_team.return_value = ["agent2", "agent4"] + mock_key.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2", "agent3"}) + ) + mock_team.return_value = RestrictedAgentAccess( + frozenset({"agent2", "agent4"}) + ) - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["agent2"] + assert result == RestrictedAgentAccess(frozenset({"agent2"})) # Case 2: Team has agents, key has none - inherit from team with patch.object( @@ -56,41 +66,124 @@ class TestAgentRequestHandler: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: - mock_key.return_value = [] - mock_team.return_value = ["team_agent1", "team_agent2"] + mock_key.return_value = UnrestrictedAgentAccess() + mock_team.return_value = RestrictedAgentAccess( + frozenset({"team_agent1", "team_agent2"}) + ) - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["team_agent1", "team_agent2"] + assert result == RestrictedAgentAccess( + frozenset({"team_agent1", "team_agent2"}) + ) - # Case 3: No restrictions - returns empty list (allow all) + # Case 3: Key has agents, team has none - key restrictions stand with patch.object( AgentRequestHandler, "_get_allowed_agents_for_key" ) as mock_key: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: - mock_key.return_value = [] - mock_team.return_value = [] + mock_key.return_value = RestrictedAgentAccess(frozenset({"key_agent1"})) + mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert result == [] + assert result == RestrictedAgentAccess(frozenset({"key_agent1"})) + + # Case 4: No grant anywhere - unrestricted (documented open-by-default) + with patch.object( + AgentRequestHandler, "_get_allowed_agents_for_key" + ) as mock_key: + with patch.object( + AgentRequestHandler, "_get_allowed_agents_for_team" + ) as mock_team: + mock_key.return_value = UnrestrictedAgentAccess() + mock_team.return_value = UnrestrictedAgentAccess() + + result = await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) + assert result == UnrestrictedAgentAccess() + + async def test_disjoint_key_and_team_grants_deny_every_agent(self): + """LIT-5143: a key restricted to one agent inside a team restricted to another + must reach nothing. The empty intersection used to read as "no restrictions", + so adding the team grant handed the key every agent on the proxy.""" + mock_user_auth: Final = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-team" + ) + + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + mock_key.return_value = RestrictedAgentAccess(frozenset({"agent-alpha"})) + mock_team.return_value = RestrictedAgentAccess(frozenset({"agent-beta"})) + + assert await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset()) + + for agent_id in ("agent-alpha", "agent-beta", "agent-secret"): + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id=agent_id, user_api_key_auth=mock_user_auth + ) + is False + ), agent_id + + async def test_empty_access_group_denies_every_agent(self): + """LIT-5143: a key restricted to an access group that resolves to no agents is + restricted to nothing, not unrestricted. A failed group lookup still fails open.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + mock_user_auth.object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="obj-1", + agents=[], + agent_access_groups=["group-with-no-agents"], + ) + + with patch.object( + AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock + ) as mock_groups: + mock_groups.return_value = [] + + assert await AgentRequestHandler._get_allowed_agents_for_key( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset()) + + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id="agent-secret", user_api_key_auth=mock_user_auth + ) + is False + ) + + with patch.object( + AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock + ) as mock_groups: + mock_groups.side_effect = Exception("DB Error") + + assert await AgentRequestHandler._get_allowed_agents_for_key( + user_api_key_auth=mock_user_auth + ) == UnrestrictedAgentAccess() async def test_is_agent_allowed_respects_permissions(self): """ - Test is_agent_allowed: returns True if agent in allowed list or if no restrictions. + Test is_agent_allowed: returns True if agent in allowed list or if unrestricted. Returns False if agent not in allowed list. """ mock_user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") # Agent in allowed list - should be allowed with patch.object( - AgentRequestHandler, "get_allowed_agents" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = ["agent1", "agent2"] + mock_get_allowed.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2"}) + ) assert ( await AgentRequestHandler.is_agent_allowed( agent_id="agent1", user_api_key_auth=mock_user_auth @@ -100,9 +193,11 @@ class TestAgentRequestHandler: # Agent not in allowed list - should be denied with patch.object( - AgentRequestHandler, "get_allowed_agents" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = ["agent1", "agent2"] + mock_get_allowed.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2"}) + ) assert ( await AgentRequestHandler.is_agent_allowed( agent_id="agent3", user_api_key_auth=mock_user_auth @@ -110,11 +205,23 @@ class TestAgentRequestHandler: is False ) - # Empty list means no restrictions - should allow any agent + # Restricted to nothing - should deny every agent with patch.object( - AgentRequestHandler, "get_allowed_agents" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = [] + mock_get_allowed.return_value = RestrictedAgentAccess(frozenset()) + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id="any_agent", user_api_key_auth=mock_user_auth + ) + is False + ) + + # Unrestricted - should allow any agent + with patch.object( + AgentRequestHandler, "resolve_agent_access" + ) as mock_get_allowed: + mock_get_allowed.return_value = UnrestrictedAgentAccess() assert ( await AgentRequestHandler.is_agent_allowed( agent_id="any_agent", user_api_key_auth=mock_user_auth @@ -126,17 +233,19 @@ class TestAgentRequestHandler: """ Test that when user_api_key_auth is None, all agents are allowed (no restrictions). """ - result = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=None) - assert result == [] + result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=None) + assert result == UnrestrictedAgentAccess() is_allowed = await AgentRequestHandler.is_agent_allowed( agent_id="any_agent", user_api_key_auth=None ) assert is_allowed is True - async def test_get_allowed_agents_handles_errors_gracefully(self): + async def test_resolve_agent_access_handles_errors_gracefully(self): """ - Test that errors during permission lookup are handled gracefully (returns empty list). + Test that errors during permission lookup are handled gracefully. This stays + fail-open for now to preserve existing availability behavior; fail-closed is + tracked separately. """ mock_user_auth = UserAPIKeyAuth( api_key="test-key", @@ -152,12 +261,12 @@ class TestAgentRequestHandler: AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: mock_key.side_effect = Exception("DB Error") - mock_team.return_value = [] + mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert result == [] + assert result == UnrestrictedAgentAccess() async def test_get_allowed_agents_for_key_via_access_group_ids(self): """ @@ -181,7 +290,9 @@ class TestAgentRequestHandler: result = await AgentRequestHandler._get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["agent-from-ag-1", "agent-from-ag-2"] + assert result == RestrictedAgentAccess( + frozenset({"agent-from-ag-1", "agent-from-ag-2"}) + ) async def test_get_allowed_agents_for_key_combines_native_and_access_groups(self): """ @@ -211,4 +322,92 @@ class TestAgentRequestHandler: result = await AgentRequestHandler._get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["agent-from-ag", "native-agent-1"] + assert result == RestrictedAgentAccess( + frozenset({"agent-from-ag", "native-agent-1"}) + ) + + async def test_is_agent_allowed_accepts_legacy_config_agent_id_grants(self): + """LIT-5144: object_permission grants stored under the pre-fix full-entry hash + must keep authorizing the agent after its id became name-based.""" + entry: Final = { + "agent_name": "granted-agent", + "agent_card_params": { + "name": "Granted Agent", + "url": "http://localhost", + "version": "1.0.0", + }, + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("granted-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + registry, + ): + with patch.object(AgentRequestHandler, "resolve_agent_access") as mock_get_allowed: + for grant, expected in ( + (RestrictedAgentAccess(frozenset({legacy_id})), True), + (RestrictedAgentAccess(frozenset({agent.agent_id})), True), + (RestrictedAgentAccess(frozenset({"unrelated-agent-id"})), False), + (RestrictedAgentAccess(frozenset()), False), + (UnrestrictedAgentAccess(), True), + ): + mock_get_allowed.return_value = grant + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id=agent.agent_id, + user_api_key_auth=mock_user_auth, + ) + is expected + ), grant + + async def test_resolve_agent_access_intersects_legacy_team_grant_with_stable_key_grant(self): + """LIT-5144: a team grant stored under the pre-fix full-entry hash and a key grant + stored under the name-based id name the same agent; the intersection must resolve + to that agent instead of collapsing to an empty set.""" + entry: Final = { + "agent_name": "shared-agent", + "agent_card_params": { + "name": "Shared Agent", + "url": "http://localhost", + "version": "1.0.0", + }, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("shared-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="test-team") + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + registry, + ): + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + for key_grant, team_grant in ( + ( + RestrictedAgentAccess(frozenset({agent.agent_id})), + RestrictedAgentAccess(frozenset({legacy_id})), + ), + ( + RestrictedAgentAccess(frozenset({legacy_id})), + RestrictedAgentAccess(frozenset({agent.agent_id})), + ), + ( + RestrictedAgentAccess(frozenset({legacy_id})), + UnrestrictedAgentAccess(), + ), + ): + mock_key.return_value = key_grant + mock_team.return_value = team_grant + assert await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset({agent.agent_id})), (key_grant, team_grant) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py index a51d6f6abcc..5193ad989dc 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_header_isolation.py @@ -227,7 +227,7 @@ async def test_each_agent_gets_only_its_own_static_headers(): # --------------------------------------------------------------------------- -# Unit tests: create_a2a_client (httpx client per call + timeout defaults) +# Unit tests: create_a2a_client (shared client left untouched + timeout defaults) # --------------------------------------------------------------------------- @@ -253,17 +253,13 @@ async def _fake_create_client(base_url, client_config=None, **kwargs): @pytest.mark.asyncio -async def test_create_a2a_client_uses_fresh_httpx_client(): +async def test_create_a2a_client_leaves_the_shared_client_untouched(): """ - Two calls to create_a2a_client with different extra_headers must produce - distinct underlying httpx clients — preventing header bleed between agents. - - The test checks: - 1. get_async_httpx_client was called twice (once per create_a2a_client call). - 2. The two returned A2A clients carry distinct httpx client objects (direct - proof of header isolation, not just cache-key difference). - 3. The cache-key param differs between calls (so the real LRU cache cannot - return the same httpx client even under load). + The client from get_async_httpx_client is shared across every A2A caller, so + create_a2a_client must neither write an agent's headers onto it nor derive its + cache key from them. Both would tie one agent's credentials to a cached object + that outlives the request. Header isolation on the wire is covered end to end in + tests/test_litellm/a2a_protocol/test_main.py. """ pytest.importorskip("a2a.client") from litellm.a2a_protocol.main import create_a2a_client @@ -281,39 +277,23 @@ async def test_create_a2a_client_uses_fresh_httpx_client(): new=AsyncMock(side_effect=_fake_create_client), ), ): - a2a_client_a = await create_a2a_client( + await create_a2a_client( base_url="http://agent-a:9999", extra_headers={"Authorization": "Bearer a"}, ) - a2a_client_b = await create_a2a_client( + await create_a2a_client( base_url="http://agent-b:9999", extra_headers={"Authorization": "Bearer b"}, ) - assert ( - len(captured_calls) == 2 - ), "create_a2a_client should call get_async_httpx_client once per invocation" + assert len(captured_calls) == 2 - # Direct proof: the two A2A clients must carry distinct httpx client objects. - # If they share one, mutating agent-B's Authorization header would bleed into A. - httpx_a = getattr(a2a_client_a, "_litellm_httpx_client", None) - httpx_b = getattr(a2a_client_b, "_litellm_httpx_client", None) - assert httpx_a is not None, "a2a_client_a missing _litellm_httpx_client" - assert httpx_b is not None, "a2a_client_b missing _litellm_httpx_client" - assert httpx_a is not httpx_b, ( - "create_a2a_client returned the same httpx client for two agents with " - "different headers — Authorization header will bleed between agents" - ) - - # Also verify the cache-key param differs so the LRU cache never conflates them. - key_a = captured_calls[0]["params"].get("disable_aiohttp_transport") - key_b = captured_calls[1]["params"].get("disable_aiohttp_transport") - assert key_a is not None, "cache-key param 'disable_aiohttp_transport' missing" - assert key_b is not None, "cache-key param 'disable_aiohttp_transport' missing" - assert key_a != key_b, ( - f"create_a2a_client used the same cache key for two agents with different " - f"headers — headers will bleed: key_a={key_a!r}, key_b={key_b!r}" - ) + for call in captured_calls: + call["client"].headers.update.assert_not_called() + assert list(call["params"]) == ["timeout"], ( + f"create_a2a_client passed extra client params {sorted(call['params'])}; " + "anything header-derived here gives every agent its own cached client" + ) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index f506535b5e1..a87f5384c6f 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -1,10 +1,14 @@ """Unit tests for AgentRegistry DB operations.""" +import hashlib +import json +from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry +from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry, GrantMigrationResult def _sample_agent_card_params() -> dict: @@ -153,7 +157,7 @@ def test_load_agents_from_db_and_config_skips_incomplete_config_entries(): registry.load_agents_from_db_and_config(db_agents=None) - assert registry.get_agent_list() == [] + assert registry.get_agent_list() == () @pytest.mark.parametrize( @@ -278,4 +282,149 @@ def test_load_agents_from_config_with_an_empty_list_clears_the_remembered_agents assert registry.config_agents == () registry.load_agents_from_db_and_config(db_agents=None) - assert registry.get_agent_list() == [], "a removed config agent must not come back on the next rebuild" + assert registry.get_agent_list() == (), "a removed config agent must not come back on the next rebuild" + + +def test_config_agent_id_survives_static_header_secret_rotation(): + """LIT-5144: the id was a hash of the whole entry, so rotating a static_headers secret silently + re-identified the agent and orphaned every grant pointing at it.""" + base_entry: Final = { + "agent_name": "rotating-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry_v1: Final = AgentRegistry() + registry_v1.load_agents_from_config([base_entry]) + agent_v1: Final = registry_v1.get_agent_by_name("rotating-agent") + assert agent_v1 is not None + + registry_v2: Final = AgentRegistry() + registry_v2.load_agents_from_config([{**base_entry, "static_headers": {"x-upstream-token": "token-v2"}}]) + agent_v2: Final = registry_v2.get_agent_by_name("rotating-agent") + assert agent_v2 is not None + + assert agent_v1.agent_id == agent_v2.agent_id + + +def test_config_agent_ids_differ_when_only_the_agent_name_differs(): + """Two entries identical except for agent_name must not collapse onto one id.""" + registry: Final = AgentRegistry() + registry.load_agents_from_config( + [ + {"agent_name": "agent-a", "agent_card_params": _sample_agent_card_params()}, + {"agent_name": "agent-b", "agent_card_params": _sample_agent_card_params()}, + ] + ) + + ids: Final = {agent.agent_id for agent in registry.get_agent_list()} + assert len(ids) == 2 + + +def test_legacy_full_entry_hash_still_resolves_the_config_agent(): + """Grants and clients created before LIT-5144 hold the old full-entry hash; it must keep resolving.""" + entry: Final = { + "agent_name": "legacy-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("legacy-agent") + assert agent is not None + + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + assert registry.config_agent_legacy_ids[legacy_id] == agent.agent_id + assert legacy_id in registry.ids_for_agent(agent.agent_id) + assert agent.agent_id in registry.ids_for_agent(agent.agent_id) + + resolved: Final = registry.get_agent_by_id(legacy_id) + assert resolved is not None + assert resolved.agent_id == agent.agent_id + assert registry.get_agent_by_id("nonexistent-id") is None + + +def test_public_agent_groups_holding_the_legacy_id_still_mark_the_config_agent_public(monkeypatch): + """LIT-5144: config.yaml written before the fix stores the full-entry hash in + public_agent_groups; the agent must stay public after its id became name-based.""" + import litellm + + entry: Final = { + "agent_name": "public-agent", + "agent_card_params": _sample_agent_card_params(), + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("public-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + assert legacy_id != agent.agent_id + + monkeypatch.setattr(litellm, "public_agent_groups", [legacy_id]) + assert [a.agent_id for a in registry.get_public_agent_list()] == [agent.agent_id] + + monkeypatch.setattr(litellm, "public_agent_groups", ["unrelated-id"]) + assert registry.get_public_agent_list() == () + + monkeypatch.setattr(litellm, "public_agent_groups", None) + assert registry.get_public_agent_list() == () + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_persists_stable_ids_into_grant_rows(): + """LIT-5144: the startup migration rewrites stored legacy full-entry hashes to the stable + name id, so a later secret rotation (which re-mints the legacy hash) cannot orphan grants.""" + entry: Final = { + "agent_name": "migrated-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + agent: Final = registry.get_agent_by_name("migrated-agent") + assert agent is not None + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + + row: Final = SimpleNamespace(object_permission_id="op-1", agents=[legacy_id, "unrelated-id", agent.agent_id]) + table: Final = MagicMock() + table.find_many = AsyncMock(return_value=[row]) + table.update_many = AsyncMock(return_value=1) + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=1, missed=0) + table.find_many.assert_awaited_once_with(where={"agents": {"has_some": (legacy_id,)}}) + table.update_many.assert_awaited_once_with( + where={"object_permission_id": "op-1", "agents": {"equals": (legacy_id, "unrelated-id", agent.agent_id)}}, + data={"agents": (agent.agent_id, "unrelated-id")}, + ) + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_reports_compare_and_swap_misses(): + """A concurrently edited row makes the CAS update affect zero rows; the result must + surface that as missed so the startup task knows to retry instead of reporting success.""" + entry: Final = { + "agent_name": "contended-agent", + "agent_card_params": _sample_agent_card_params(), + "static_headers": {"x-upstream-token": "token-v1"}, + } + registry: Final = AgentRegistry() + registry.load_agents_from_config([entry]) + legacy_id: Final = hashlib.sha256(json.dumps(entry, sort_keys=True).encode()).hexdigest() + + row: Final = SimpleNamespace(object_permission_id="op-1", agents=[legacy_id]) + table: Final = MagicMock() + table.find_many = AsyncMock(return_value=[row]) + table.update_many = AsyncMock(return_value=0) + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=0, missed=1) + + +@pytest.mark.asyncio +async def test_migrate_legacy_grant_ids_no_ops_without_config_agents(): + """Without config agents there are no legacy hashes to translate, so the DB is never queried.""" + registry: Final = AgentRegistry() + table: Final = MagicMock() + table.find_many = AsyncMock() + + assert await registry.migrate_legacy_grant_ids(table=table) == GrantMigrationResult(rewritten=0, missed=0) + table.find_many.assert_not_awaited() diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 3e097711ad7..ea196bda529 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -6,6 +6,10 @@ from fastapi.testclient import TestClient from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) from litellm.proxy.agent_endpoints.endpoints import ( _attach_keys_to_agents, _check_agent_management_permission, @@ -308,7 +312,7 @@ async def test_attach_keys_to_agents_groups_by_agent_and_omits_secret(): # Query is scoped to the agents being returned, not the whole key table. where = mock_prisma.db.litellm_verificationtoken.find_many.call_args.kwargs["where"] - assert where == {"agent_id": {"in": ["agent-1", "agent-2"]}} + assert where == {"agent_id": {"in": ("agent-1", "agent-2")}} # agent-1 gets both of its keys; agent-2 gets None. assert agent_without_keys.keys is None @@ -503,11 +507,14 @@ class TestAgentRBACProxyAdminViewOnly: ] self.mock_registry = MagicMock() self.mock_registry.get_agent_list = MagicMock(return_value=self.agents) + self.mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry) - self.allowed_agents_spy = AsyncMock(return_value=["someone-elses-agent"]) + self.allowed_agents_spy = AsyncMock( + return_value=RestrictedAgentAccess(frozenset({"someone-elses-agent"})) + ) monkeypatch.setattr( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", self.allowed_agents_spy, ) @@ -536,9 +543,9 @@ class TestAgentRBACProxyAdminViewOnly: self.allowed_agents_spy.assert_awaited_once() def test_should_still_redact_secrets_for_view_only_admin(self): - """An unrestricted viewer (empty allowlist means no restrictions) sees the - same agents as an admin but with keys stripped and litellm_params masked.""" - self.allowed_agents_spy.return_value = [] + """An unrestricted viewer sees the same agents as an admin but with keys + stripped and litellm_params masked.""" + self.allowed_agents_spy.return_value = UnrestrictedAgentAccess() viewer_resp = self._list_agents(self.viewer_client) admin_resp = self._list_agents(self.admin_client) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py index 86eb7a83079..ccf5942c89d 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py @@ -13,6 +13,9 @@ from unittest.mock import AsyncMock, Mock, patch import pytest +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, +) from litellm.proxy.agent_endpoints.model_list_helpers import ( append_agents_to_model_group, append_agents_to_model_info, @@ -37,14 +40,14 @@ async def test_append_agents_to_model_group(): ) # Mock AgentRequestHandler at its source location - mock_get_allowed_agents = AsyncMock(return_value=["test-agent-id"]) + mock_get_allowed_agents = AsyncMock(return_value=RestrictedAgentAccess(frozenset({"test-agent-id"}))) # Mock global_agent_registry mock_registry = Mock() mock_registry.get_agent_by_id = Mock(return_value=mock_agent) with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", mock_get_allowed_agents, ): with patch( @@ -80,14 +83,14 @@ async def test_append_agents_to_model_info(): ) # Mock AgentRequestHandler at its source location - mock_get_allowed_agents = AsyncMock(return_value=["agent-123"]) + mock_get_allowed_agents = AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-123"}))) # Mock global_agent_registry mock_registry = Mock() mock_registry.get_agent_by_id = Mock(return_value=mock_agent) with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", mock_get_allowed_agents, ): with patch( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index a5211ba83e7..298f8a31b64 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -5105,6 +5105,80 @@ async def test_common_checks_personal_user_budget_skipped_for_team_key(): assert result is True +@pytest.mark.asyncio +async def test_common_checks_personal_user_budget_enforced_on_team_key_when_flag_enabled(): + """general_settings.apply_user_budget_to_team_keys opts a deployment into + charging the key owner's personal budget on team-scoped keys too. + + Same fixture as the default-off test above, so a regression that ignores the + flag lets this call through instead of raising. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) + team = LiteLLM_TeamTable(team_id="t1", spend=0.0, max_budget=1000.0) + token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 if counter_key == "spend:user:u1" else 0.0 + + async def _no_membership(*args, **kwargs): + return None + + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + ), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=team, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={"apply_user_budget_to_team_keys": True}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + assert "ExceededBudget: User=u1" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_common_checks_personal_user_budget_still_enforced_on_personal_key_with_flag_enabled(): + """The flag only widens enforcement to team keys; personal keys keep blocking.""" + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0) + token = UserAPIKeyAuth(token="k1", user_id="u1") + + async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs): + return 999.0 if counter_key == "spend:user:u1" else 0.0 + + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.proxy_server.get_current_spend", _spend_by_counter + ): + with pytest.raises(litellm.BudgetExceededError): + await common_checks( + request_body={"messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=user, + end_user_object=None, + global_proxy_spend=None, + general_settings={"apply_user_budget_to_team_keys": True}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=token, + request=MagicMock(spec=Request), + ) + + @pytest.mark.parametrize( "scope, route, expect_blocked", [ @@ -5464,6 +5538,63 @@ async def test_get_project_object_db_fetch_returns_cached_obj(): assert result.project_alias == "proj" +@pytest.mark.asyncio +async def test_project_allowlist_enforced_when_key_models_empty(): + """ + LIT-3803: a project-bound key with models=[] has no key-level restriction, + but the project allowlist must still 403 team models outside it. + """ + from litellm.proxy._types import ( + LiteLLM_ProjectTableCachedObj, + ProxyErrorTypes, + ProxyException, + ) + from litellm.proxy.auth.auth_checks import _run_project_checks, can_key_call_model + + valid_token = UserAPIKeyAuth( + api_key="hashed-key", + project_id="p-1", + team_id="t-1", + models=[], + ) + project = LiteLLM_ProjectTableCachedObj( + project_id="p-1", + team_id="t-1", + models=["gemini-2.5-flash-image", "gemini-3.1-flash-lite-preview"], + ) + + assert ( + await can_key_call_model( + model="gemini-2.5-flash", + llm_model_list=None, + valid_token=valid_token, + llm_router=None, + ) + is True + ) + + await _run_project_checks( + project_object=project, + _model="gemini-2.5-flash-image", + llm_router=None, + skip_budget_checks=True, + valid_token=valid_token, + proxy_logging_obj=MagicMock(), + ) + + with pytest.raises(ProxyException) as exc_info: + await _run_project_checks( + project_object=project, + _model="gemini-2.5-flash", + llm_router=None, + skip_budget_checks=True, + valid_token=valid_token, + proxy_logging_obj=MagicMock(), + ) + assert exc_info.value.type == ProxyErrorTypes.project_model_access_denied + assert exc_info.value.code == "403" + + def test_is_user_proxy_admin_rejects_view_only_admin(): """This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an Admin Viewer answering True here would gain every write route. Read parity for diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index f56b8e113a9..e6c0eaee3c4 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -735,3 +735,28 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): litellm.vertex_language_models.discard(fake_model) litellm.add_known_models(model_cost_map={}) assert fake_model not in litellm.models_by_provider["vertex_ai"] + +def test_get_complete_model_list_drops_no_default_models_sentinel(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=["no-default-models", "model-a"], + team_models=[], + proxy_model_list=["model-a", "model-b"], + user_model=None, + infer_model_from_keys=False, + ) + assert result == ["model-a"] + + +def test_get_complete_model_list_sentinel_only_grants_nothing(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=["no-default-models"], + team_models=["no-default-models"], + proxy_model_list=["model-a", "model-b"], + user_model=None, + infer_model_from_keys=False, + ) + assert result == [] diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index bc5bb877bd0..60d9689dc0b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -218,6 +218,46 @@ async def test_fail_closed_budget_enforcement_reaches_reservation( ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "general_settings,expected_flag", + [ + ({"apply_user_budget_to_team_keys": True}, True), + ({"apply_user_budget_to_team_keys": False}, False), + ({}, False), + ], +) +async def test_apply_user_budget_to_team_keys_reaches_reservation( + general_settings, expected_flag +): + """The opt-in lives in general_settings but is consumed inside + _get_budget_counters, so it has to be threaded through reserve_budget_for_request + or the reservation path keeps exempting team keys while the read path enforces.""" + user_api_key_auth_obj = UserAPIKeyAuth(token="test_token") + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request", + new=AsyncMock(return_value=None), + ) as mock_reserve: + await _reserve_budget_after_common_checks( + user_api_key_auth_obj=user_api_key_auth_obj, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + llm_router=None, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + skip_budget_checks=False, + general_settings=general_settings, + ) + + assert ( + mock_reserve.await_args.kwargs["apply_user_budget_to_team_keys"] is expected_flag + ) + + @pytest.mark.asyncio async def test_should_not_reuse_cached_key_object_for_request_state(): key_cache = DualCache() @@ -4746,6 +4786,117 @@ async def test_user_api_key_auth_does_not_overwrite_end_user_id_set_by_builder() setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +async def test_user_api_key_auth_authenticates_before_raising_malformed_body_error(): + """Regression (LIT-4780): a body that fails to parse must still be authenticated + first, so the rejected request's trace carries the caller's key / team / user + identity instead of an anonymous root span. The parse error is re-raised + unchanged once identity is seeded.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + builder_token = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="team-1") + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = b'{}{"model": "gpt-4o"}' + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + return_value=builder_token, + ) as mock_builder, + patch( + "litellm.proxy.auth.user_api_key_auth._run_centralized_common_checks", + new_callable=AsyncMock, + ) as mock_common_checks, + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + patch( + "litellm.proxy.auth.user_api_key_auth.seed_request_identity", + ) as mock_seed, + ): + with pytest.raises(ProxyException) as exc_info: + await user_api_key_auth(request=request, api_key="Bearer sk-test") + + assert "Invalid JSON payload" in str(exc_info.value.message) + assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST) + mock_builder.assert_awaited_once() + assert mock_seed.call_args.args[0] is builder_token + # authorization must not run for a request that is about to be rejected: + # ``common_checks`` reserves budget against live spend counters that only the + # endpoint's post-call path releases, and the endpoint never runs here + mock_common_checks.assert_not_awaited() + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + +@pytest.mark.asyncio +async def test_user_api_key_auth_malformed_body_with_rejected_key_still_returns_the_parse_error(): + """The body is read before the key is authenticated, so a caller who sends both a + malformed body and a key that fails auth gets the 400. Authenticating the request + first (LIT-4780) must not turn that into the auth status code.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + request = Request( + scope={ + "type": "http", + "headers": [(b"content-type", b"application/json")], + "method": "POST", + } + ) + request._url = URL(url="/chat/completions") + request._body = b'{}{"model": "gpt-4o"}' + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth._user_api_key_auth_builder", + new_callable=AsyncMock, + side_effect=ProxyException( + message="Authentication Error, invalid key", + type="auth_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ), + ), + patch( + "litellm.proxy.auth.user_api_key_auth.RouteChecks.should_call_route", + ), + ): + with pytest.raises(ProxyException) as exc_info: + await user_api_key_auth(request=request, api_key="Bearer sk-bad") + + assert "Invalid JSON payload" in str(exc_info.value.message) + assert exc_info.value.code == str(status.HTTP_400_BAD_REQUEST) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + def _proxy_attrs_for_db_lookup(): """Minimal proxy_server attributes for driving the real ``_user_api_key_auth_builder`` down to the DB key lookup.""" diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index e758aa5ca7f..f9193db143e 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -2253,3 +2253,176 @@ async def test_cancel__provider_only_resolves_named_vertex_credentials(cancel_ha "vertex_location": "us-central1", "vertex_credentials": "/creds/customer-sa.json", } + + +# =========================================================================== # +# require_managed_files - raw provider ids must not reach the provider. # +# # +# Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to # +# these routes is forwarded under the shared provider credentials with no # +# tenant check, so any caller who learns another tenant's id can read its # +# batch, reuse its file as batch input, or cancel its job. These lock the # +# guard on every batches route that accepts a caller-supplied id. # +# =========================================================================== # + + +def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-provider-id") -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _unified_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "managed-id", "gpt-4o-mini", "file-provider-id", "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: bool = True + object_access: bool = True + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.file_access + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.object_access + + +@pytest.mark.asyncio +async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness): + set_body( + harness, + { + "input_file_id": "file-victim-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "400" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__model_encoded_input_file_id_rejected_when_managed_files_required(harness): + """A model-encoded id is client-forgeable and has no ownership row, so it is + not a managed file id and must be rejected like any other raw id.""" + set_body( + harness, + { + "input_file_id": AZURE_FILE_ID, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "400" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_create__raw_input_file_id_allowed_when_managed_files_not_required(harness): + set_body( + harness, + { + "input_file_id": "file-victim-abc123", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + + with patch.object(litellm, "require_managed_files", False): + await call_create(harness) + + assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123" + + +@pytest.mark.asyncio +async def test_create__other_teams_unified_input_file_id_rejected(harness): + set_body( + harness, + { + "input_file_id": _unified_file_id(), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + ) + harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_create(harness) + + assert exc.value.code == "403" + harness.litellm_acreate.assert_not_called() + harness.router_acreate.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_retrieve(retrieve_harness, "batch-victim-abc123") + + assert exc.value.code == "400" + retrieve_harness.litellm_aretrieve.assert_not_called() + retrieve_harness.router_aretrieve.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness): + retrieve_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await call_retrieve(retrieve_harness, _unified_batch_id()) + + assert retrieve_harness.router_aretrieve.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_harness): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await call_cancel(cancel_harness, "batch-victim-abc123") + + assert exc.value.code == "400" + cancel_harness.litellm_acancel.assert_not_called() + cancel_harness.router_acancel.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness): + cancel_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await call_cancel(cancel_harness, _unified_batch_id()) + + assert cancel_harness.router_acancel.call_count == 1 diff --git a/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py new file mode 100644 index 00000000000..468e8aabae8 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -0,0 +1,161 @@ +import asyncio +import json +from typing import Iterable, List, Optional, Tuple +from unittest.mock import patch + +import pytest +from redis.asyncio import Redis + +from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( + AUTH_CACHE_INVALIDATION_CHANNEL, + AuthCacheInvalidationSubscriber, + publish_auth_cache_invalidation, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + +class _RecordingRedisClient(Redis): + def __init__(self) -> None: + self.published: List[Tuple[str, str]] = [] + + async def publish(self, channel: str, message: str) -> int: + self.published.append((channel, message)) + return 1 + + +class _FailingPublishRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + raise ConnectionError("redis down") + + +class _QueuePubSub: + def __init__(self, initial_messages: Iterable[object] = ()) -> None: + self.queue: "asyncio.Queue[object]" = asyncio.Queue() + for message in initial_messages: + self.queue.put_nowait(message) + self.subscribed_channels: List[str] = [] + self.closed = False + + async def subscribe(self, *channels: str) -> None: + self.subscribed_channels.extend(channels) + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]: + try: + return await asyncio.wait_for(self.queue.get(), timeout) + except asyncio.TimeoutError: + return None + + async def aclose(self) -> None: + self.closed = True + + +class _ScriptedPubSubRedisClient(Redis): + def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None: + self._scripted_pubsubs = iter(pubsubs) + + def pubsub(self) -> _QueuePubSub: + return next(self._scripted_pubsubs) + + +class _FakeRedisCache: + def __init__(self, client: object, namespace: Optional[str] = None) -> None: + self._client = client + self.namespace = namespace + + def init_async_client(self) -> object: + return self._client + + +def _invalidation_message(cache_key: str) -> dict: + return {"type": "message", "data": json.dumps({"cache_key": cache_key}).encode()} + + +@pytest.mark.asyncio +async def test_publish_sends_cache_key_json_on_channel() -> None: + client = _RecordingRedisClient() + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + assert client.published == [(AUTH_CACHE_INVALIDATION_CHANNEL, json.dumps({"cache_key": "project_id:p-1"}))] + + +@pytest.mark.asyncio +async def test_publish_uses_namespaced_channel() -> None: + client = _RecordingRedisClient() + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client, namespace="ns1"), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + assert client.published[0][0] == f"ns1:{AUTH_CACHE_INVALIDATION_CHANNEL}" + + +@pytest.mark.asyncio +async def test_publish_noops_without_coordination_redis() -> None: + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=None, + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + +@pytest.mark.asyncio +async def test_publish_swallows_redis_errors() -> None: + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=_FailingPublishRedisClient()), + ): + await publish_auth_cache_invalidation(cache_key="project_id:p-1") + + +@pytest.mark.asyncio +async def test_subscriber_deletes_local_cache_entry_on_message() -> None: + """ + The cross-worker half of LIT-3803: a worker that did not handle the project + mutation must drop its in-memory copy when the invalidation broadcast lands, + instead of serving the stale object until the TTL expires. + """ + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + assert cache.in_memory_cache.get_cache("project_id:p-1") is not None + + pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + ) + subscriber.start() + try: + for _ in range(200): + if cache.in_memory_cache.get_cache("project_id:p-1") is None: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert cache.in_memory_cache.get_cache("project_id:p-1") is None + assert pubsub.subscribed_channels == [AUTH_CACHE_INVALIDATION_CHANNEL] + + +@pytest.mark.asyncio +async def test_subscriber_ignores_malformed_messages() -> None: + cache = UserApiKeyCache() + cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) + + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), + user_api_key_cache=cache, + ) + subscriber._apply_message({"type": "message", "data": b"not json"}) + subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()}) + subscriber._apply_message("raw string") + subscriber._apply_message(None) + + assert cache.in_memory_cache.get_cache("project_id:p-1") is not None diff --git a/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py new file mode 100644 index 00000000000..e2fa1de6962 --- /dev/null +++ b/tests/test_litellm/proxy/credential_endpoints/test_endpoints.py @@ -0,0 +1,91 @@ +"""Tests for the credential management endpoints.""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.proxy_server import app +from litellm.types.utils import CredentialItem + +client = TestClient(app) + + +def _as_admin(): + return UserAPIKeyAuth(api_key="test-key", user_role="proxy_admin") + + +def _patch_credential(name: str, body: dict): + missing = object() + previous_override = app.dependency_overrides.get(user_api_key_auth, missing) + app.dependency_overrides[user_api_key_auth] = _as_admin + try: + return client.patch( + f"/credentials/{name}", + json=body, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + if previous_override is missing: + app.dependency_overrides.pop(user_api_key_auth, None) + else: + app.dependency_overrides[user_api_key_auth] = previous_override + + +def test_update_credential_answers_404_when_the_credential_does_not_exist(): + """Regression: the handler used to ``return handle_exception_on_proxy(e)``, which makes + the exception the response body and lets FastAPI answer 200, so a write the handler + rejected read as a success to every caller that checks the status. The dashboard's API + client branches on the status, so it reported a failed edit as applied.""" + with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( + "litellm.proxy.credential_endpoints.endpoints.CredentialsRepository" + ) as repository: + repository.return_value.find_by_name = AsyncMock(return_value=None) + + response = _patch_credential( + "definitely-not-there", + {"credential_name": "definitely-not-there", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, + ) + + assert response.status_code == 404, f"rejected write answered {response.status_code}: {response.text}" + assert "error" in response.json() + + +def test_update_credential_answers_500_when_the_database_is_not_connected(): + """The other rejection this handler raises must carry its own status too.""" + with patch("litellm.proxy.proxy_server.prisma_client", None): + response = _patch_credential( + "any-name", + {"credential_name": "any-name", "credential_values": {"api_key": "sk-x"}, "credential_info": {}}, + ) + + assert response.status_code == 500, f"rejected write answered {response.status_code}: {response.text}" + + +def test_update_credential_still_answers_200_on_a_successful_write(): + """The fix must not turn a legitimate update into an error; the dashboard and the + Playwright credentials spec both assert the success path.""" + stored = CredentialItem( + credential_name="existing", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch( + "litellm.proxy.proxy_server.master_key", "sk-test-master" + ), patch("litellm.proxy.credential_endpoints.endpoints.CredentialsRepository") as repository: + repository.return_value.find_by_name = AsyncMock(return_value=stored) + repository.return_value.update_by_name = AsyncMock(return_value=None) + + response = _patch_credential( + "existing", + {"credential_name": "existing", "credential_values": {"api_key": "sk-new"}, "credential_info": {}}, + ) + + assert response.status_code == 200, response.text + assert response.json()["success"] is True diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index aa7d01bc880..0df11f224a2 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -93,6 +93,19 @@ class TestBuildTransaction: def test_requests_without_a_routing_decision_are_skipped(self, metadata: dict): assert _build(metadata=metadata) is None + def test_the_tier_the_decision_recorded_is_carried_onto_the_transaction(self): + transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": "reasoning"})) + assert transaction is not None and transaction.tier == "reasoning" + + @pytest.mark.parametrize("tier", [None, "", 3, {"tier": "medium"}]) + def test_a_decision_without_a_usable_tier_records_no_tier(self, tier: object): + transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": tier})) + assert transaction is not None and transaction.tier is None + + def test_a_decision_that_never_mentions_tier_records_no_tier(self): + transaction = _build() + assert transaction is not None and transaction.tier is None + def test_router_name_falls_back_to_the_payload_model_group(self): transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"})) assert transaction is not None and transaction.router_name == "live-auto" @@ -167,7 +180,11 @@ class _FakeClient: self.db = _FakeDB(failures, poison_session) -def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, 0, 0)) -> AutoRouterTurnTransaction: +def _transaction( + session_id: str = "s1", + at: datetime = datetime(2026, 8, 1, 12, 0, 0), + tier: str | None = "medium", +) -> AutoRouterTurnTransaction: return AutoRouterTurnTransaction( api_key="k1", session_id=session_id, @@ -182,6 +199,7 @@ def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, cache_hit=False, cache_ttl_seconds=None, cache_touched=False, + tier=tier, ) @@ -201,7 +219,7 @@ class TestFlush: assert sql == UPSERT_AUTOROUTER_SESSION_SQL assert params == ( "k1", "s1", "live-auto", "complexity", "bedrock/haiku", - "2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0, + "2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0, "medium", ) def test_a_connect_error_retries_the_same_statement(self): diff --git a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py index 18a238b7545..c8e0338eeaa 100644 --- a/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py +++ b/tests/test_litellm/proxy/db/test_prisma_planned_engine_restart.py @@ -6,13 +6,24 @@ subprocess), and the engine-death watcher / in-flight transport-error retries must not treat that planned restart as a crash and recreate the client a second time. +A planned restart is also invisible to the database itself, so a ``SELECT 1`` +health probe that races the kill/connect window fails with a transport error +against the engine's local HTTP port. That failure must not be reported as a +``db_exceptions`` DB failure, or every IAM refresh cycle raises a false alarm. + Symbols pinned here: - ``PrismaWrapper._expected_engine_deaths`` - ``PrismaWrapper._engine_generation`` + - ``PrismaWrapper.engine_generation`` + - ``PrismaWrapper.wait_for_planned_engine_replacement`` - ``PrismaWrapper.on_engine_replaced`` - ``PrismaWrapper.recreate_prisma_client`` (expected_generation guard) - ``PrismaWrapper._safe_refresh_token`` (refresh coalescing) - ``RoutingPrismaWrapper.recreate_prisma_client`` (guard forwarding) + - ``PrismaClient.health_check`` (planned-replacement alert suppression) + - ``PrismaClient._probe_target_wrapper`` + - ``PrismaClient._probe_answers_now`` + - ``PrismaClient._planned_engine_replacement_absorbed`` """ import asyncio @@ -21,22 +32,29 @@ import signal import sys import urllib.parse from datetime import datetime, timedelta +from typing import Any, List from unittest.mock import AsyncMock, MagicMock, call, patch +import httpx import pytest from prisma import Prisma as GeneratedPrisma +from prisma.engine.errors import EngineConnectionError sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path from litellm.proxy.db.prisma_client import PrismaWrapper +from litellm.proxy.utils import PrismaClient @pytest.fixture(autouse=True) def mock_prisma_binary(): """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests.""" mock_module = MagicMock() + # Production code isinstance-checks against this, which a bare MagicMock + # attribute cannot satisfy. + mock_module.engine.errors.EngineConnectionError = EngineConnectionError with patch.dict(sys.modules, {"prisma": mock_module}): yield mock_module @@ -49,6 +67,87 @@ def _make_wrapper(engine_pid: int = 111, iam: bool = False) -> PrismaWrapper: return PrismaWrapper(original_prisma=prisma, iam_token_db_auth=iam) +def _make_prisma_client(db: Any) -> PrismaClient: + """A ``PrismaClient`` whose ``db`` is a real wrapper and whose alerting + hook is observable.""" + proxy_logging_obj = MagicMock() + proxy_logging_obj.failure_handler = AsyncMock() + client = PrismaClient( + database_url="postgresql://user:pass@localhost:5432/db", + proxy_logging_obj=proxy_logging_obj, + ) + client.db = db + client._db_watchdog_reconnect_timeout_seconds = 5.0 + return client + + +_real_asyncio_sleep = asyncio.sleep + + +async def _yield_to_loop(times: int = 10) -> None: + """Let already-scheduled tasks make progress. + + Bound to the real ``asyncio.sleep`` at import time: the tests below patch + ``asyncio.sleep`` to skip the SIGTERM/SIGKILL grace, and an ``AsyncMock`` + stand-in never yields to the event loop, which would silently leave every + background task un-started and the assertions vacuous. + """ + for _ in range(times): + await _real_asyncio_sleep(0) + + +async def _await_health_check_reports() -> None: + """Await the fire-and-forget reporting tasks ``health_check()`` scheduled. + + Selected by coroutine qualname rather than by draining every pending task, + so an unrelated background task can never make these assertions pass by + accident. + """ + reports = [ + task + for task in asyncio.all_tasks() + if getattr(task.get_coro(), "__qualname__", "") + == "PrismaClient._report_health_check_failure" + ] + if reports: + await asyncio.gather(*reports, return_exceptions=True) + + +def _fails_then_answers(error: Exception, failures: int = 3) -> Any: + """Raise ``error`` for the first ``failures`` probes, then answer. + + ``health_check`` retries up to three times, so this exhausts the retries and + still lets the confirmation probe that decides suppression succeed. Without + that, a test would report for the wrong reason: the confirmation probe would + fail too, masking whether the error type was classified at all. + """ + seen: List[int] = [] + + async def _query_raw(_sql: str) -> Any: + seen.append(1) + if len(seen) <= failures: + raise error + return [{"?column?": 1}] + + return _query_raw + + +def _blocking_replacement(gate: asyncio.Event, fail: bool = False) -> MagicMock: + """A replacement Prisma whose ``connect()`` parks until ``gate`` is set. + + Holds ``_reconnection_lock`` open for as long as the test needs, which is + how a health probe is made to fail *while* a planned replacement is in + flight rather than after it. + """ + + async def _connect(*_: Any, **__: Any) -> None: + await gate.wait() + if fail: + raise ConnectionRefusedError("database is down") + + return MagicMock(connect=AsyncMock(side_effect=_connect)) + + def _token_db_url(created: datetime, expires_in: int = 900) -> str: """Build a DATABASE_URL whose password is a parseable RDS IAM token.""" token = ( @@ -603,3 +702,348 @@ async def test_recreate_caps_expected_engine_deaths_set(mock_prisma_binary): await wrapper.recreate_prisma_client("postgresql://new") assert wrapper._expected_engine_deaths == {111} + + +@pytest.mark.asyncio +async def test_wait_for_planned_engine_replacement_returns_once_recreate_settles( + mock_prisma_binary, +): + wrapper = _make_wrapper(engine_pid=111) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + waiter = asyncio.create_task(wrapper.wait_for_planned_engine_replacement(5.0)) + await _yield_to_loop() + blocked_while_in_flight = not waiter.done() + + gate.set() + await recreate + await waiter + + assert { + "blocked_while_in_flight": blocked_while_in_flight, + "generation": wrapper.engine_generation, + } == {"blocked_while_in_flight": True, "generation": 1} + + +@pytest.mark.asyncio +async def test_wait_for_planned_engine_replacement_gives_up_at_timeout( + mock_prisma_binary, +): + """A replacement that never settles must not stall the caller forever; the + caller then sees an unchanged generation and reports the failure.""" + wrapper = _make_wrapper(engine_pid=111) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + await asyncio.wait_for( + wrapper.wait_for_planned_engine_replacement(0.05), timeout=5.0 + ) + gave_up_with_replacement_still_in_flight = not recreate.done() + + gate.set() + await recreate + + assert gave_up_with_replacement_still_in_flight is True + + +@pytest.mark.asyncio +async def test_health_check_does_not_alert_when_probe_races_a_completed_replacement( + mock_prisma_binary, +): + """The reported bug: an IAM-refresh engine recreate makes a concurrent + readiness probe fail transiently, and that failure was alerting as a DB + exception on every refresh cycle. + + The reporting task is drained while the replacement is still in flight, + which is when it runs in production; a decision taken at that instant sees + an engine generation that has not moved yet. + """ + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + probe_result = await client.health_check() + + drain = asyncio.create_task(_await_health_check_reports()) + await _yield_to_loop() + alerts_while_replacement_in_flight = ( + client.proxy_logging_obj.failure_handler.await_count + ) + + gate.set() + await recreate + await drain + + assert { + "probe_result": probe_result, + "probe_attempts": wrapper.query_raw.await_count, + "alerts_while_in_flight": alerts_while_replacement_in_flight, + "alerts": client.proxy_logging_obj.failure_handler.await_count, + } == { + "probe_result": [{"?column?": 1}], + "probe_attempts": 3, + "alerts_while_in_flight": 0, + "alerts": 0, + } + + +@pytest.mark.asyncio +async def test_health_check_alerts_when_a_completed_replacement_still_cannot_reach_the_database( + mock_prisma_binary, +): + """``Prisma.connect()`` polls the query engine's own ``/status`` endpoint + rather than round-tripping to the database, so a replacement can complete + against a database that is still unreachable. The engine generation alone + must not be enough to stay silent.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=httpx.ConnectError("All connection attempts failed") + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + gate.set() + await recreate + await _await_health_check_reports() + + assert { + "replacement_completed": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"replacement_completed": 1, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_check_alerts_when_the_replacement_never_completes( + mock_prisma_binary, +): + """A real outage also has a replacement in flight, but it fails, so the + engine generation never advances and the probe failure must still alert.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=httpx.ConnectError("All connection attempts failed") + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate, fail=True) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(httpx.ConnectError): + await client.health_check() + + gate.set() + with pytest.raises(ConnectionRefusedError): + await recreate + await _await_health_check_reports() + + call_types: List[str] = [ + c.kwargs["call_type"] + for c in client.proxy_logging_obj.failure_handler.await_args_list + ] + assert { + "generation": wrapper.engine_generation, + "alerted": len(call_types) > 0, + "call_types": set(call_types), + } == {"generation": 0, "alerted": True, "call_types": {"health_check"}} + + +@pytest.mark.asyncio +async def test_health_check_alerts_for_non_connection_errors_during_a_replacement( + mock_prisma_binary, +): + """Suppression is scoped to transport failures. A query the database itself + rejected is a real defect and must alert even mid-replacement, and even + though the database is plainly reachable a moment later.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock(side_effect=_fails_then_answers(ValueError("malformed SELECT"))) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + wrapper.recreate_prisma_client("postgresql://new") + ) + await _yield_to_loop() + assert wrapper._reconnection_lock.locked() is True + + with pytest.raises(ValueError): + await client.health_check() + + gate.set() + await recreate + await _await_health_check_reports() + + assert { + "generation": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"generation": 1, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_check_alerts_for_a_transient_failure_with_no_engine_replacement( + mock_prisma_binary, +): + """Suppression is scoped to failures a planned replacement explains. A + transport blip that self-heals with no engine replacement at all still + alerts, so the gate cannot be widened into silencing every failure whose + database happens to answer a moment later.""" + wrapper = _make_wrapper(engine_pid=111) + client = _make_prisma_client(wrapper) + wrapper.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + + probe_result = await client.health_check() + await _await_health_check_reports() + + assert { + "probe_result": probe_result, + "generation": wrapper.engine_generation, + "alerted": client.proxy_logging_obj.failure_handler.await_count > 0, + } == {"probe_result": [{"?column?": 1}], "generation": 0, "alerted": True} + + +@pytest.mark.asyncio +async def test_health_probe_stays_on_its_target_when_reader_availability_flips( + mock_prisma_binary, monkeypatch +): + """Routing is re-resolved on every attribute access, so a reader that + recovers mid-call would otherwise send the probe to a different engine than + the one whose generation is being checked, and blame the wrong replacement. + The probe follows the wrapper it was handed.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://reader") + writer = _make_wrapper(engine_pid=111) + reader = _make_wrapper(engine_pid=222) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + client = _make_prisma_client(routing) + writer.query_raw = AsyncMock(return_value=[{"writer": 1}]) + reader.query_raw = AsyncMock(return_value=[{"reader": 1}]) + + routing._reader_unavailable = False + target = client._probe_target_wrapper() + routing._reader_unavailable = True + result = await client._run_health_probe(target) + + assert { + "target_is_reader": target is reader, + "result": result, + "reader_probes": reader.query_raw.await_count, + "writer_probes": writer.query_raw.await_count, + } == { + "target_is_reader": True, + "result": [{"reader": 1}], + "reader_probes": 1, + "writer_probes": 0, + } + + +@pytest.mark.asyncio +async def test_health_check_consults_the_reader_wrapper_under_read_replica_routing( + mock_prisma_binary, monkeypatch +): + """``query_raw`` is routed to the reader, so a reader-side planned + replacement is the one that explains a probe failure.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", "postgresql://reader") + writer = _make_wrapper(engine_pid=111) + reader = _make_wrapper(engine_pid=222) + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + client = _make_prisma_client(routing) + reader.query_raw = AsyncMock( + side_effect=[ + httpx.ConnectError("All connection attempts failed"), + [{"?column?": 1}], + [{"?column?": 1}], + ] + ) + gate = asyncio.Event() + mock_prisma_binary.Prisma.return_value = _blocking_replacement(gate) + + with patch("os.kill"), patch("asyncio.sleep", new_callable=AsyncMock): + recreate = asyncio.create_task( + reader.recreate_prisma_client("postgresql://new-reader") + ) + await _yield_to_loop() + assert reader._reconnection_lock.locked() is True + + probe_result = await client.health_check() + + drain = asyncio.create_task(_await_health_check_reports()) + await _yield_to_loop() + alerts_while_replacement_in_flight = ( + client.proxy_logging_obj.failure_handler.await_count + ) + + gate.set() + await recreate + await drain + + assert { + "probe_target_is_the_reader": client._probe_target_wrapper() is reader, + "writer_generation": writer.engine_generation, + "reader_generation": reader.engine_generation, + "probe_result": probe_result, + "alerts_while_in_flight": alerts_while_replacement_in_flight, + "alerts": client.proxy_logging_obj.failure_handler.await_count, + } == { + "probe_target_is_the_reader": True, + "writer_generation": 0, + "reader_generation": 1, + "probe_result": [{"?column?": 1}], + "alerts_while_in_flight": 0, + "alerts": 0, + } diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py b/tests/test_litellm/proxy/fine_tuning_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py new file mode 100644 index 00000000000..b54787bf428 --- /dev/null +++ b/tests/test_litellm/proxy/fine_tuning_endpoints/test_endpoints.py @@ -0,0 +1,278 @@ +""" +require_managed_files enforcement for litellm/proxy/fine_tuning_endpoints/endpoints.py + +Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to these +routes is forwarded to the provider under the shared proxy credentials with no tenant +check, so any caller who learns another tenant's file id can train on it, and any +caller who learns another tenant's job id can read or cancel it. + +Each test asserts BOTH that the request is rejected AND that every downstream provider +seam stayed untouched, so a guard that raises after the provider call would still fail. +""" + +import base64 +import os +import sys +from contextlib import ExitStack +from dataclasses import dataclass +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import Response + +import litellm +import litellm.proxy.fine_tuning_endpoints.endpoints as endpoints +import litellm.proxy.proxy_server as proxy_server +from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.types.llms.openai import LiteLLMFineTuningJobCreate +from litellm.types.utils import LiteLLMFineTuningJob, SpecialEnums + +RAW_FILE_ID = "file-victim-abc123" +RAW_JOB_ID = "ftjob-victim-abc123" + + +def _unified_file_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _unified_job_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_GENERIC_RESPONSE_COMPLETE_STR.value.format("gpt-4o-mini-id", RAW_JOB_ID) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +def _job() -> LiteLLMFineTuningJob: + job = LiteLLMFineTuningJob( + id=RAW_JOB_ID, + created_at=1234567890, + fine_tuned_model=None, + finished_at=None, + hyperparameters={"n_epochs": 1}, + model="gpt-4o-mini", + object="fine_tuning.job", + organization_id="org-test", + result_files=[], + seed=0, + status="running", + trained_tokens=None, + training_file=RAW_FILE_ID, + validation_file=None, + ) + job._hidden_params = {} + return job + + +class FakeRequest: + def __init__(self): + self.headers = {} + self.query_params = {} + + async def json(self): + return {} + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: bool = True + object_access: bool = True + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.file_access + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return self.object_access + + +class Seams: + def __init__(self, router: MagicMock, litellm_calls: dict[str, AsyncMock], logging: MagicMock): + self.router = router + self.litellm_calls = litellm_calls + self.logging = logging + + def assert_no_provider_call(self) -> None: + for name, mock in self.litellm_calls.items(): + assert mock.call_count == 0, f"litellm.{name} was called" + for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job"): + assert getattr(self.router, name).call_count == 0, f"router.{name} was called" + + +@pytest.fixture +def seams(): + logging = MagicMock(spec=ProxyLogging) + logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"]) + logging.post_call_failure_hook = AsyncMock() + logging.update_request_status = AsyncMock() + logging.get_proxy_hook = MagicMock(return_value=None) + + router = MagicMock(spec=Router) + router.acreate_fine_tuning_job = AsyncMock(return_value=_job()) + router.aretrieve_fine_tuning_job = AsyncMock(return_value=_job()) + router.acancel_fine_tuning_job = AsyncMock(return_value=_job()) + + litellm_calls = { + name: AsyncMock(return_value=_job()) + for name in ("acreate_fine_tuning_job", "aretrieve_fine_tuning_job", "acancel_fine_tuning_job") + } + + with ExitStack() as stack: + stack.enter_context( + patch.object( + ProxyBaseLLMRequestProcessing, + "common_processing_pre_call_logic", + AsyncMock(side_effect=lambda self=None, **kw: (self.data if self else {}, MagicMock())), + ) + ) + stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", MagicMock(return_value={}))) + for name, mock in litellm_calls.items(): + stack.enter_context(patch.object(litellm, name, mock)) + stack.enter_context(patch.object(proxy_server, "llm_router", router)) + stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging)) + stack.enter_context(patch.object(proxy_server, "premium_user", True)) + stack.enter_context(patch.object(proxy_server, "general_settings", {})) + stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock())) + stack.enter_context(patch.object(proxy_server, "version", "test-version")) + stack.enter_context(patch.object(endpoints, "fine_tuning_config", [{"custom_llm_provider": "openai"}])) + yield Seams(router=router, litellm_calls=litellm_calls, logging=logging) + + +async def _create(training_file: str, validation_file: str | None = None): + return await endpoints.create_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_request=LiteLLMFineTuningJobCreate( + model="gpt-4o-mini", + training_file=training_file, + validation_file=validation_file, + custom_llm_provider="openai", + ), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +async def _retrieve(job_id: str): + return await endpoints.retrieve_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_job_id=job_id, + custom_llm_provider="openai", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +async def _cancel(job_id: str): + return await endpoints.cancel_fine_tuning_job( + request=FakeRequest(), + fastapi_response=Response(), + fine_tuning_job_id=job_id, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + +@pytest.mark.asyncio +async def test_create__raw_training_file_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(RAW_FILE_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__raw_validation_file_rejected_when_managed_files_required(seams): + """The validation file is uploaded and readable exactly like the training file, + so a managed training_file must not smuggle a raw validation_file past the guard.""" + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(_unified_file_id(), validation_file=RAW_FILE_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__unified_training_file_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await _create(_unified_file_id()) + + assert seams.router.acreate_fine_tuning_job.call_count == 1 + + +@pytest.mark.asyncio +async def test_create__other_teams_unified_training_file_rejected(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False) + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _create(_unified_file_id()) + + assert exc.value.code == "403" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_create__raw_training_file_allowed_when_managed_files_not_required(seams): + with patch.object(litellm, "require_managed_files", False): + await _create(RAW_FILE_ID) + + assert seams.litellm_calls["acreate_fine_tuning_job"].call_count == 1 + + +@pytest.mark.asyncio +async def test_retrieve__raw_job_id_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _retrieve(RAW_JOB_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_retrieve__unified_job_id_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await _retrieve(_unified_job_id()) + + assert seams.router.aretrieve_fine_tuning_job.call_count == 1 + + +@pytest.mark.asyncio +async def test_cancel__raw_job_id_rejected_when_managed_files_required(seams): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(ProxyException) as exc: + await _cancel(RAW_JOB_ID) + + assert exc.value.code == "400" + seams.assert_no_provider_call() + + +@pytest.mark.asyncio +async def test_cancel__unified_job_id_allowed_when_managed_files_required(seams): + seams.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub() + + with patch.object(litellm, "require_managed_files", True): + await _cancel(_unified_job_id()) + + assert seams.router.acancel_fine_tuning_job.call_count == 1 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 76a695ce3fd..837fb93d331 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -7,6 +7,7 @@ import os import sys from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException @@ -15,12 +16,18 @@ sys.path.insert(0, os.path.abspath("../../../../../..")) import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockContentChunkResult, BedrockGuardrail, _redact_pii_matches, ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockContentItem, + BedrockTextContent, +) from litellm.types.utils import CallTypes, ModelResponse @@ -53,9 +60,7 @@ async def test__redact_pii_matches_function(): redacted_response = _redact_pii_matches(response_with_pii) # Verify that PII matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" @@ -173,12 +178,8 @@ async def test__redact_pii_matches_multiple_assessments(): redacted_response = _redact_pii_matches(response_multiple_assessments) # Verify all PII in all assessments are redacted - assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ - "piiEntities" - ] + assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"]["piiEntities"] assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" @@ -199,9 +200,7 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII mock_bedrock_response = MagicMock() @@ -239,20 +238,11 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" - ) as mock_debug, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug") as mock_debug, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method that should log the redacted response @@ -275,37 +265,23 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): bedrock_response_log_call = call break - assert ( - bedrock_response_log_call is not None - ), "Should have logged Bedrock AI response" + assert bedrock_response_log_call is not None, "Should have logged Bedrock AI response" # Extract the logged response data - logged_response = bedrock_response_log_call[0][ - 1 - ] # Second argument to debug call + logged_response = bedrock_response_log_call[0][1] # Second argument to debug call # Verify that the logged response has redacted PII assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] - == "[REDACTED]" + logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) # Verify other fields are preserved assert logged_response["action"] == "GUARDRAIL_INTERVENED" - assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["type"] - == "PHONE" - ) + assert logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["type"] == "PHONE" slg_list = request_data["metadata"]["standard_logging_guardrail_information"] assert ( - slg_list[0]["guardrail_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + slg_list[0]["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) @@ -319,9 +295,7 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII original_response_data = { @@ -361,17 +335,10 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method @@ -385,19 +352,12 @@ async def test_bedrock_guardrail_original_response_not_modified(): # (The json() method should return the original data) original_data = mock_bedrock_response.json() assert ( - original_data["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + original_data["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" ) # Verify that the returned BedrockGuardrailResponse contains original data - assert ( - result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "+1 412 555 1212" - ) + assert result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" print("Original response not modified test passed") @@ -454,18 +414,14 @@ async def test__redact_pii_matches_preserves_non_pii_entities(): redacted_response = _redact_pii_matches(response_with_mixed_data) # Verify that PII entity matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" # Verify that regex matches are also redacted (updated behavior) - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" @@ -496,9 +452,7 @@ async def test_pii_redaction_matches_debug_output_format(): "assessments": [ { "invocationMetrics": { - "guardrailCoverage": { - "textCharacters": {"guarded": 84, "total": 84} - }, + "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, "guardrailProcessingLatency": 322, "usage": { "contentPolicyImageUnits": 0, @@ -553,9 +507,7 @@ async def test_pii_redaction_matches_debug_output_format(): redacted_response = _redact_pii_matches(original_response) # Verify the redacted response matches your expected debug output - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] # All PII matches should be redacted assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" @@ -570,34 +522,19 @@ async def test_pii_redaction_matches_debug_output_format(): assert pii_entities[0]["detected"] == True # Verify that the original response is unchanged - original_pii_entities = original_response["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"] - assert ( - original_pii_entities[0]["match"] == "John Smith" - ), "Original should be unchanged" - assert ( - original_pii_entities[1]["match"] == "324-12-3212" - ), "Original should be unchanged" - assert ( - original_pii_entities[2]["match"] == "607-456-7890" - ), "Original should be unchanged" + original_pii_entities = original_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert original_pii_entities[0]["match"] == "John Smith", "Original should be unchanged" + assert original_pii_entities[1]["match"] == "324-12-3212", "Original should be unchanged" + assert original_pii_entities[2]["match"] == "607-456-7890", "Original should be unchanged" # Verify all other metadata is preserved in redacted response assert redacted_response["action"] == "GUARDRAIL_INTERVENED" assert redacted_response["actionReason"] == "Guardrail blocked." assert redacted_response["blockedResponse"] == "Input blocked by PII policy" - assert ( - redacted_response["assessments"][0]["invocationMetrics"][ - "guardrailProcessingLatency" - ] - == 322 - ) + assert redacted_response["assessments"][0]["invocationMetrics"]["guardrailProcessingLatency"] == 322 print("PII redaction matches debug output format test passed") - print( - f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" - ) + print(f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}") print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") @@ -632,14 +569,10 @@ async def test__redact_pii_matches_with_regex_matches(): redacted_response = _redact_pii_matches(response_with_regex) # Verify that regex matches are redacted - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" - assert ( - regexes[1]["match"] == "[REDACTED]" - ), "Credit card regex match should be redacted" + assert regexes[1]["match"] == "[REDACTED]", "Credit card regex match should be redacted" # Verify other fields are preserved assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" @@ -648,13 +581,9 @@ async def test__redact_pii_matches_with_regex_matches(): assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" # Verify original response is unchanged - original_regexes = response_with_regex["assessments"][0][ - "sensitiveInformationPolicy" - ]["regexes"] + original_regexes = response_with_regex["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" - assert ( - original_regexes[1]["match"] == "4111-1111-1111-1111" - ), "Original should be unchanged" + assert original_regexes[1]["match"] == "4111-1111-1111-1111", "Original should be unchanged" print("Regex matches redaction test passed") @@ -690,31 +619,17 @@ async def test__redact_pii_matches_with_custom_words(): # Verify that custom word matches are redacted custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "First custom word match should be redacted" - assert ( - custom_words[1]["match"] == "[REDACTED]" - ), "Second custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "First custom word match should be redacted" + assert custom_words[1]["match"] == "[REDACTED]", "Second custom word match should be redacted" # Verify other fields are preserved - assert ( - custom_words[0]["action"] == "BLOCKED" - ), "Custom word action should be preserved" - assert ( - custom_words[1]["action"] == "ANONYMIZED" - ), "Custom word action should be preserved" + assert custom_words[0]["action"] == "BLOCKED", "Custom word action should be preserved" + assert custom_words[1]["action"] == "ANONYMIZED", "Custom word action should be preserved" # Verify original response is unchanged - original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ - "customWords" - ] - assert ( - original_custom_words[0]["match"] == "confidential_data" - ), "Original should be unchanged" - assert ( - original_custom_words[1]["match"] == "secret_information" - ), "Original should be unchanged" + original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"]["customWords"] + assert original_custom_words[0]["match"] == "confidential_data", "Original should be unchanged" + assert original_custom_words[1]["match"] == "secret_information", "Original should be unchanged" print("Custom words redaction test passed") @@ -750,41 +665,21 @@ async def test__redact_pii_matches_with_managed_words(): redacted_response = _redact_pii_matches(response_with_managed_words) # Verify that managed word matches are redacted - managed_words = redacted_response["assessments"][0]["wordPolicy"][ - "managedWordLists" - ] + managed_words = redacted_response["assessments"][0]["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "First managed word match should be redacted" - assert ( - managed_words[1]["match"] == "[REDACTED]" - ), "Second managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "First managed word match should be redacted" + assert managed_words[1]["match"] == "[REDACTED]", "Second managed word match should be redacted" # Verify other fields are preserved - assert ( - managed_words[0]["action"] == "BLOCKED" - ), "Managed word action should be preserved" - assert ( - managed_words[0]["type"] == "PROFANITY" - ), "Managed word type should be preserved" - assert ( - managed_words[1]["action"] == "ANONYMIZED" - ), "Managed word action should be preserved" - assert ( - managed_words[1]["type"] == "HATE_SPEECH" - ), "Managed word type should be preserved" + assert managed_words[0]["action"] == "BLOCKED", "Managed word action should be preserved" + assert managed_words[0]["type"] == "PROFANITY", "Managed word type should be preserved" + assert managed_words[1]["action"] == "ANONYMIZED", "Managed word action should be preserved" + assert managed_words[1]["type"] == "HATE_SPEECH", "Managed word type should be preserved" # Verify original response is unchanged - original_managed_words = response_with_managed_words["assessments"][0][ - "wordPolicy" - ]["managedWordLists"] - assert ( - original_managed_words[0]["match"] == "inappropriate_word" - ), "Original should be unchanged" - assert ( - original_managed_words[1]["match"] == "offensive_term" - ), "Original should be unchanged" + original_managed_words = response_with_managed_words["assessments"][0]["wordPolicy"]["managedWordLists"] + assert original_managed_words[0]["match"] == "inappropriate_word", "Original should be unchanged" + assert original_managed_words[1]["match"] == "offensive_term", "Original should be unchanged" print("Managed words redaction test passed") @@ -841,9 +736,7 @@ async def test__redact_pii_matches_comprehensive_coverage(): # PII entities pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] - assert ( - pii_entities[0]["match"] == "[REDACTED]" - ), "PII entity match should be redacted" + assert pii_entities[0]["match"] == "[REDACTED]", "PII entity match should be redacted" # Regex matches regexes = assessment["sensitiveInformationPolicy"]["regexes"] @@ -851,15 +744,11 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Custom words custom_words = assessment["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "Custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "Custom word match should be redacted" # Managed words managed_words = assessment["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "Managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "Managed word match should be redacted" # Verify all other fields are preserved assert pii_entities[0]["type"] == "EMAIL" @@ -868,21 +757,10 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Verify original response is unchanged original_assessment = comprehensive_response["assessments"][0] - assert ( - original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] - == "user@example.com" - ) - assert ( - original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] - == "555-123-4567" - ) - assert ( - original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" - ) - assert ( - original_assessment["wordPolicy"]["managedWordLists"][0]["match"] - == "inappropriate" - ) + assert original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "user@example.com" + assert original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] == "555-123-4567" + assert original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" + assert original_assessment["wordPolicy"]["managedWordLists"][0]["match"] == "inappropriate" print("Comprehensive coverage redaction test passed") @@ -914,9 +792,7 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method to avoid actual AWS credential loading - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -926,10 +802,12 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + ) print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") @@ -944,9 +822,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) # Create guardrail without explicit aws_bedrock_runtime_endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -960,9 +836,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -972,10 +846,10 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint from environment is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, f"Expected URL to contain env endpoint. Got: {prepped_request.url}" print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") @@ -988,9 +862,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) # Create guardrail without any custom endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -1004,9 +876,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey aws_region_name = "us-west-2" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1017,9 +887,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey # Verify that the default endpoint is used expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected default URL. Got: {prepped_request.url}" + assert prepped_request.url == expected_url, f"Expected default URL. Got: {prepped_request.url}" print(f"Default endpoint test passed. URL: {prepped_request.url}") @@ -1057,9 +925,7 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1069,10 +935,12 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch ) # Verify that the parameter takes precedence over environment variable - expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + expected_url = ( + f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + ) print(f"Parameter precedence test passed. URL: {prepped_request.url}") @@ -1081,14 +949,10 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" # Create a BedrockGuardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the make_bedrock_api_request method - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api_request: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api_request: # Test the apply_guardrail method with tool_calls in response inputs = { "texts": [], @@ -1115,14 +979,9 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): assert guardrailed_inputs is not None assert "tool_calls" in guardrailed_inputs assert len(guardrailed_inputs["tool_calls"]) == 1 - assert ( - guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" - ) + assert guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" - assert ( - guardrailed_inputs["tool_calls"][0]["function"]["arguments"] - == '{"location":"São Paulo"}' - ) + assert guardrailed_inputs["tool_calls"][0]["function"]["arguments"] == '{"location":"São Paulo"}' # Verify that the Bedrock API was NOT called since there's no text to process mock_api_request.assert_not_called() print("✅ apply_guardrail with tool_calls test passed - no API call made") @@ -1136,14 +995,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): policies (e.g. PII on model output) then returned action=NONE for non-streaming completions that go through unified_guardrail -> process_output_response. """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1168,14 +1023,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): @pytest.mark.asyncio async def test_bedrock_apply_guardrail_request_uses_INPUT_source(): """input_type='request' must call Bedrock with source=INPUT and user messages.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1258,12 +1109,8 @@ async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): # Mock AWS-related methods with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -1431,9 +1278,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" def _create_guardrail(self) -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") @pytest.mark.asyncio async def test_should_handle_all_null_policy_sub_lists(self): @@ -1554,9 +1399,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: { "sensitiveInformationPolicy": { "piiEntities": None, - "regexes": [ - {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} - ], + "regexes": [{"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"}], }, } ], @@ -1611,18 +1454,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_none_texts_in_inputs(self): """inputs[\"texts\"] is explicitly None — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {"texts": None} # Explicit None mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1645,18 +1484,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_missing_texts_key(self): """inputs has no \"texts\" key at all — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {} # No "texts" key mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1677,9 +1512,7 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Test 1: ANONYMIZED action should NOT raise exception anonymized_response = { @@ -1700,9 +1533,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): ], } - should_raise = guardrail._should_raise_guardrail_blocked_exception( - anonymized_response - ) + should_raise = guardrail._should_raise_guardrail_blocked_exception(anonymized_response) assert should_raise is False, "ANONYMIZED actions should not raise exceptions" # Test 2: BLOCKED action should raise exception @@ -1710,13 +1541,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "I can't provide that information."}], "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } + {"topicPolicy": {"topics": [{"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"}]}} ], } @@ -1738,19 +1563,13 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): } ] }, - "topicPolicy": { - "topics": [ - {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"}]}, } ], } should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert ( - should_raise is True - ), "Mixed actions with any BLOCKED should raise exceptions" + assert should_raise is True, "Mixed actions with any BLOCKED should raise exceptions" # Test 4: NONE action should not raise exception none_response = { @@ -1782,9 +1601,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): When logging_event_type is set, it must be forwarded to standard guardrail logging. When omitted, INPUT maps to pre_call (legacy). """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1795,13 +1612,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): mock_bedrock_response.json.return_value = { "action": "NONE", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ], } @@ -1811,12 +1622,8 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -1831,15 +1638,13 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): request_data=request_data, logging_event_type=GuardrailEventHooks.during_call, ) - assert ( - mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call - ) + assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call # Raw Bedrock JSON is forwarded; redaction runs once in # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. assert ( - mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] == "GG" ) @@ -1855,9 +1660,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): @pytest.mark.asyncio async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1873,15 +1676,9 @@ async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): prepared_request.headers = {} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), - patch.object( - guardrail, "_prepare_request", return_value=prepared_request - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=prepared_request) as mock_prepare_request, patch.object( guardrail, "get_guardrail_dynamic_request_body_params", @@ -1933,9 +1730,7 @@ async def test_during_call_hook_invokes_bedrock_async_moderation_hook(): "model": "gpt-4", "messages": [{"role": "user", "content": "test"}], }, - user_api_key_dict=UserAPIKeyAuth( - api_key="test_key", user_id="test_user" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), call_type="completion", ) finally: @@ -1991,11 +1786,7 @@ def test_extract_blocked_assessments_multiple_policies(): "action": "GUARDRAIL_INTERVENED", "assessments": [ { - "topicPolicy": { - "topics": [ - {"name": "Investment", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Investment", "type": "DENY", "action": "BLOCKED"}]}, "contentPolicy": { "filters": [ { @@ -2006,9 +1797,7 @@ def test_extract_blocked_assessments_multiple_policies(): } ] }, - "wordPolicy": { - "customWords": [{"match": "forbidden", "action": "BLOCKED"}] - }, + "wordPolicy": {"customWords": [{"match": "forbidden", "action": "BLOCKED"}]}, } ], } @@ -2023,13 +1812,7 @@ def test_extract_blocked_assessments_only_anonymized_returns_empty(): response = { "action": "GUARDRAIL_INTERVENED", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } assert g._extract_blocked_assessments(response) == [] @@ -2049,23 +1832,14 @@ def test_get_http_exception_includes_assessments_and_identifier(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "BLOCKED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "BLOCKED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) assert isinstance(exc, HTTPException) assert exc.status_code == 400 assert exc.detail["error"] == "Violated guardrail policy" - assert ( - exc.detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question." - ) + assert exc.detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question." assert exc.detail["guardrailIdentifier"] == "amgllac6xf3r" assert exc.detail["guardrailVersion"] == "1" assert exc.detail["assessments"][0]["policy"] == "sensitiveInformationPolicy" @@ -2088,15 +1862,11 @@ def test_extract_violation_category_names_mixed_policies(): {"name": "Tax Advice", "action": "BLOCKED"}, ] }, - "contentPolicy": { - "filters": [{"type": "VIOLENCE", "action": "BLOCKED"}] - }, + "contentPolicy": {"filters": [{"type": "VIOLENCE", "action": "BLOCKED"}]}, "wordPolicy": { "managedWordLists": [{"type": "PROFANITY", "action": "BLOCKED"}], }, - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}] - }, + "sensitiveInformationPolicy": {"piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}]}, } ], } @@ -2120,13 +1890,9 @@ def test_extract_violation_category_names_does_not_leak_user_input(): "assessments": [ { "wordPolicy": { - "customWords": [ - {"match": "secret-codeword-abc-123", "action": "BLOCKED"} - ], - }, - "sensitiveInformationPolicy": { - "regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}] + "customWords": [{"match": "secret-codeword-abc-123", "action": "BLOCKED"}], }, + "sensitiveInformationPolicy": {"regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}]}, } ], } @@ -2166,13 +1932,7 @@ def test_extract_violation_category_names_skips_anonymized(): g = _make_guardrail() response = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}] - } - } - ], + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}]}}], } assert g._extract_violation_category_names(response) == [] @@ -2190,9 +1950,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): ``tracing_detail`` so downstream loggers (OTEL, ...) can surface the raw provider verdict as a queryable attribute without re-parsing the redacted guardrail_response blob.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2202,13 +1960,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}]}}], } request_data = { @@ -2217,12 +1969,8 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2253,9 +2001,7 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): """If the Bedrock response omits ``action`` (older / partial payloads), the field must be left off ``tracing_detail`` rather than written as ``None`` — downstream code expects strings or absence, not nulls.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2266,12 +2012,8 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): mock_bedrock_response.json.return_value = {"assessments": []} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2300,13 +2042,7 @@ def test_get_http_exception_no_blocked_assessments_omits_field(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "blocked"}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) @@ -2370,9 +2106,7 @@ async def test_streaming_post_call_only_runs_output_scan(): yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: out = [] async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), @@ -2382,18 +2116,11 @@ async def test_streaming_post_call_only_runs_output_scan(): out.append(chunk) assert len(out) >= 1 - output_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT" - ] + output_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT"] assert len(output_calls) == 1 assert output_calls[0].kwargs.get("request_data") is request_data - assert ( - output_calls[0].kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) - input_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT" - ] + assert output_calls[0].kwargs.get("logging_event_type") == GuardrailEventHooks.post_call + input_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT"] assert len(input_calls) == 0 @@ -2432,9 +2159,7 @@ async def test_streaming_post_call_output_only_path_passes_request_data_to_make_ yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: async for _ in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), response=mock_stream(), @@ -2488,9 +2213,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): ) minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=UserAPIKeyAuth(), @@ -2499,10 +2222,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): sources = [c.kwargs.get("source") for c in mock_make.call_args_list] assert sources == ["OUTPUT"] - assert ( - mock_make.call_args.kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) + assert mock_make.call_args.kwargs.get("logging_event_type") == GuardrailEventHooks.post_call # --------------------------------------------------------------------------- @@ -2522,9 +2242,7 @@ _GROUNDING_RESPONSE_TEXT = "The capital of Japan is Tokyo." def _grounding_guardrail() -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") def _grounding_messages() -> list: @@ -2556,27 +2274,19 @@ def _model_response(content: str) -> ModelResponse: # Expected OUTPUT content blocks, keyed by their grounding qualifier, so the # per-test assertions read as the block sequence they expect. -_GROUNDING_SOURCE_BLOCK = { - "text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]} -} +_GROUNDING_SOURCE_BLOCK = {"text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]}} _QUERY_BLOCK = {"text": {"text": _GROUNDING_QUERY_TEXT, "qualifiers": ["query"]}} -_GUARD_BLOCK = { - "text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]} -} +_GUARD_BLOCK = {"text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]}} def _input_request(messages: list) -> dict: """Arrange a guardrail and act: build the Bedrock INPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="INPUT", messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) def _output_request(messages: list, response=None) -> dict: """Arrange a guardrail and act: build the Bedrock OUTPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="OUTPUT", response=response, messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) def test_grounding_input_strips_grounding_and_query_qualifiers(): @@ -2600,9 +2310,7 @@ def test_grounding_input_leaves_existing_guarded_text_unqualified(): """An existing guarded_text input block keeps its legacy unqualified payload.""" expected_request = {"source": "INPUT", "content": [{"text": {"text": "policy"}}]} - actual_request = _input_request( - [{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}] - ) + actual_request = _input_request([{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}]) assert actual_request == expected_request @@ -2615,9 +2323,7 @@ def test_grounding_output_assembles_source_query_and_response(): "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], } - actual_request = _output_request( - _grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(_grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2629,9 +2335,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): "content": [{"text": {"text": "Hi there."}}], } - actual_request = _output_request( - [{"role": "user", "content": "hello"}], _model_response("Hi there.") - ) + actual_request = _output_request([{"role": "user", "content": "hello"}], _model_response("Hi there.")) assert actual_request == expected_request @@ -2639,9 +2343,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): def test_grounding_output_combines_multiple_sources(): """Every grounding_source block is emitted; Bedrock combines them into one corpus.""" uk_source_text = "London is the capital of UK." - uk_source_block = { - "text": {"text": uk_source_text, "qualifiers": ["grounding_source"]} - } + uk_source_block = {"text": {"text": uk_source_text, "qualifiers": ["grounding_source"]}} messages = [ { "role": "system", @@ -2662,9 +2364,7 @@ def test_grounding_output_combines_multiple_sources(): ], } - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2709,9 +2409,7 @@ def test_grounding_source_trusted_only_from_app_roles(role, is_trusted): if is_trusted: expected_content = [_GROUNDING_SOURCE_BLOCK, *expected_content] - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == {"source": "OUTPUT", "content": expected_content} @@ -2748,12 +2446,8 @@ async def test_grounding_output_blocked_raises_400(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -2791,13 +2485,7 @@ def _blocked_bedrock_httpx_response() -> MagicMock: response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}]}}], } return response @@ -2820,12 +2508,8 @@ async def test_make_bedrock_api_request_block_raises_modify_response_when_flag_s mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2857,12 +2541,8 @@ async def test_make_bedrock_api_request_block_raises_http_400_when_flag_unset(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2903,12 +2583,8 @@ async def test_async_pre_call_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2953,12 +2629,8 @@ async def test_async_moderation_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2998,12 +2670,8 @@ async def test_async_post_call_success_hook_attaches_original_response_on_block( mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -3032,9 +2700,7 @@ async def test_apply_guardrail_propagates_modify_response_on_block(): disable_exception_on_block=True, ) - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.side_effect = ModifyResponseException( message="Sorry, the model cannot answer this question.", model="bedrock-nova-micro", @@ -3276,6 +2942,1160 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n assert response is not None +def _too_large_validation_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.json.return_value = { + "message": "Input is too long. Content size exceeds the maximum input size in text units." + } + response.text = json.dumps(response.json.return_value) + return response + + +def _other_validation_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.json.return_value = {"message": "guardrailIdentifier is not valid"} + response.text = json.dumps(response.json.return_value) + return response + + +def _throttling_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 429 + response.json.return_value = {"message": "Rate exceeded"} + response.text = json.dumps(response.json.return_value) + response.headers = {} + return response + + +def _too_large_throttling_httpx_response() -> MagicMock: + """The shape AWS actually returns for an oversized ApplyGuardrail request when + the guardrail has an active content-filter policy: a 429 ThrottlingException, + not the documented 400 ValidationException. Message taken from a live call.""" + response = MagicMock() + response.status_code = 429 + response.json.return_value = { + "message": ( + "Input text size (3273 text units) exceeds the maximum allowed " + "(1000 text units) for the content filter policy (Classic tier)." + ) + } + response.text = json.dumps(response.json.return_value) + response.headers = {} + return response + + +def _passing_bedrock_httpx_response(marker: str) -> MagicMock: + """A successful ApplyGuardrail response tagged with `marker` so tests can + verify which chunk produced which output/usage after merging.""" + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "NONE", + "outputs": [{"text": marker}], + "assessments": [], + "usage": {"contentPolicyUnits": 1}, + } + return response + + +def _blocking_bedrock_httpx_response(marker: str) -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": marker}], + "assessments": [{"topicPolicy": {"topics": [{"name": marker, "type": "DENY", "action": "BLOCKED"}]}}], + "usage": {"contentPolicyUnits": 1}, + } + return response + + +def _bedrock_guardrail_for_chunk_tests() -> "BedrockGuardrail": + return BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunks_on_too_large_validation_error(): + """A too-large 400 on the whole-content call must trigger a bisect-and-retry, + and the two chunk responses must be merged (assessments concatenated, usage + summed, outputs concatenated) rather than losing either half's result.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "first half of a very long message"}, + {"role": "user", "content": "second half of a very long message"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _blocking_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + detail = exc_info.value.detail + assert exc_info.value.status_code == 400 + assert "chunk-2" in detail["bedrock_guardrail_response"] + assert detail["assessments"][0]["matches"][0]["name"] == "chunk-2" + + +@pytest.mark.asyncio +async def test_apply_guardrail_merges_usage_and_outputs_across_chunks_when_both_pass(): + """When both chunks pass clean, the merged response must still carry both + chunks' outputs/usage forward (needed for accurate logging/telemetry) and + must not itself raise.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-1", "chunk-2"] + assert result.get("usage", {}).get("contentPolicyUnits") == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_recurses_past_first_bisection_into_four_chunks(): + """A payload that is still too large after one bisection must keep splitting + -- chunking is not capped at two pieces. Four messages where both the + whole-content call AND both first-level halves are too large must recurse + one level deeper into four chunks that all fit, not give up after the + first split.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "message one"}, + {"role": "user", "content": "message two"}, + {"role": "user", "content": "message three"}, + {"role": "user", "content": "message four"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + responses = [ + _too_large_validation_httpx_response(), # whole content: [1,2,3,4] + _too_large_validation_httpx_response(), # first half: [1,2] + _passing_bedrock_httpx_response("message one"), + _passing_bedrock_httpx_response("message two"), + _too_large_validation_httpx_response(), # second half: [3,4] + _passing_bedrock_httpx_response("message three"), + _passing_bedrock_httpx_response("message four"), + ] + + async def _post_side_effect(*_args, **_kwargs): + return responses.pop(0) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 7 + assert not responses + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["message one", "message two", "message three", "message four"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_does_not_chunk_when_grounding_present(): + """Contextual-grounding requests are scored holistically against the whole + source; chunking them would silently produce misleading grounding scores. + A too-large error on a grounded request must propagate unchanged, not be + bisected.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + { + "role": "system", + "content": [{"type": "grounding_source", "text": "reference source text"}], + }, + {"role": "user", "content": "what does the source say?"}, + ] + model_response = ModelResponse() + model_response.choices = [litellm.Choices(message=litellm.Message(content="a grounded answer", role="assistant"))] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="OUTPUT", + messages=messages, + response=model_response, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_does_not_chunk_on_non_size_validation_error(): + """A 400 for an unrelated validation problem (e.g. a bad guardrail id) must + not trigger chunking -- retrying a bad-config error split into pieces would + just fail twice more and mask the real problem.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _other_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + assert "not valid" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_on_unsplittable_text_propagates_original_error(): + """A too-large error on content that has been bisected down to text too + short to split further (< 2 characters) must propagate the original error + rather than looping or crashing.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "a"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_on_single_item_splits_by_text_and_succeeds(): + """A too-large error on content that is already down to a single content + item must be bisected by that item's own text (not abandoned), so an + oversized single message can still be scanned successfully in halves.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "one giant single block of text"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response(f"half-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + response = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 3 + assert response.get("action") == "NONE" + + +def _raised_bedrock_error(status_code: int, message: str) -> httpx.HTTPStatusError: + """A non-200 the way `AsyncHTTPHandler.post` actually surfaces it. + + That handler calls `response.raise_for_status()`, so in production a non-200 from + Bedrock arrives as a raised `httpx.HTTPStatusError` carrying the response, never + as a returned response object. Tests that return the response instead exercise a + branch real traffic never reaches. A real `httpx.Response` is used rather than a + MagicMock because the transport helper branches on + `isinstance(err_response, httpx.Response)`.""" + response = httpx.Response( + status_code=status_code, + json={"message": message}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/guardrail"), + ) + return httpx.HTTPStatusError(message, request=response.request, response=response) + + +_TOO_LARGE_MESSAGE = "Input is too long. Content size exceeds the maximum input size in text units." + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunking_logs_once_when_client_raises_for_status(): + """The too-large attempt recovered by chunking must still produce exactly one + telemetry entry when the HTTP client raises for status, which is what really + happens: `AsyncHTTPHandler.post` calls `raise_for_status()`. + + Regression for per-attempt `guardrail_failed_to_respond` entries leaking out of + the transport helper on a request that ultimately succeeded, which made a + recovered request look like several failures plus a success.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise _raised_bedrock_error(400, _TOO_LARGE_MESSAGE) + return _passing_bedrock_httpx_response(f"chunk-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + assert result.get("action") == "NONE" + statuses = [call.kwargs.get("guardrail_status") for call in mock_log.call_args_list] + assert statuses == ["success"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_unrecoverable_failure_still_logs_once_when_client_raises(): + """Suppressing the transport helper's per-attempt logging must not swallow the only + record of a genuine failure: an unsplittable too-large request still has to produce + exactly one `guardrail_failed_to_respond` entry, not zero.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "x"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + async def _post_side_effect(*_args, **_kwargs): + raise _raised_bedrock_error(400, _TOO_LARGE_MESSAGE) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + statuses = [call.kwargs.get("guardrail_status") for call in mock_log.call_args_list] + assert statuses == ["guardrail_failed_to_respond"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_single_item_split_twice_still_yields_one_output_per_item(): + """One oversized content item that needs two levels of text bisection ends up + as four text fragments, and all four must still collapse back into exactly + ONE output entry, because they all came from one original content item. + + Downstream masking (`_apply_masking_to_messages`) walks the merged outputs by + a running index across the original, unchunked message list, so emitting more + than one entry for a single message shifts every later message's masked text + onto the wrong message and drops the surplus. Regression for fragment + grouping assuming fragments only ever arrive as adjacent sibling *pairs*, + which holds for one bisection level but not for two.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "aaaa bbbb cccc dddd eeee ffff gggg hhhh"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + responses = [ + _too_large_validation_httpx_response(), # whole single item + _too_large_validation_httpx_response(), # first half + _passing_bedrock_httpx_response("q1"), + _passing_bedrock_httpx_response("q2"), + _too_large_validation_httpx_response(), # second half + _passing_bedrock_httpx_response("q3"), + _passing_bedrock_httpx_response("q4"), + ] + + async def _post_side_effect(*_args, **_kwargs): + return responses.pop(0) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 7 + assert not responses + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["q1q2q3q4"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunk_retries_after_throttling_then_succeeds(): + """A chunk call throttled with a 429 must be retried with backoff and + eventually succeed, rather than surfacing the 429 to the caller.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _throttling_httpx_response() + if call_count == 3: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.asyncio.sleep", + new_callable=AsyncMock, + ) as mock_sleep, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 4 + mock_sleep.assert_awaited() + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-1", "chunk-2"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunking_logs_exactly_once_as_success(): + """A too-large 400 that is recovered by chunking must not leave behind a + 'guardrail_failed_to_respond' telemetry entry for the initial oversized + attempt: the whole logical request (1 too-large attempt + 2 chunk + attempts here) must produce exactly one standard-logging entry, and it + must reflect the eventual success, not the transient too-large failure.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "success" + + +@pytest.mark.asyncio +async def test_apply_guardrail_unrecoverable_failure_logs_exactly_once_as_failed(): + """A too-large error that cannot be recovered (chunking disabled by + contextual grounding) must still log exactly once, as a failure -- not be + silently dropped by the chunking telemetry consolidation.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + { + "role": "system", + "content": [{"type": "grounding_source", "text": "reference source text"}], + }, + {"role": "user", "content": "what does the source say?"}, + ] + model_response = ModelResponse() + model_response.choices = [litellm.Choices(message=litellm.Message(content="a grounded answer", role="assistant"))] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="OUTPUT", + messages=messages, + response=model_response, + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_post.assert_awaited_once() + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunk_merge_preserves_masking_position(): + """An earlier chunk that comes back clean (empty `outputs`) must not + shift a later chunk's masked text onto the wrong message. Regression for: + flattening outputs without positional metadata let a later chunk's PII + redaction get applied to the first message while the actual PII-bearing + message (in a later chunk) was forwarded unmasked.""" + guardrail = BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + mask_request_content=True, + ) + + request_data = { + "model": "bedrock-nova-micro", + "messages": [ + {"role": "user", "content": "clean chunk with nothing to mask"}, + {"role": "user", "content": "chunk with PII: John Doe"}, + ], + } + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + def _clean_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"action": "NONE", "assessments": []} + return response + + def _masked_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "chunk with PII: [NAME]"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [{"type": "NAME", "match": "John Doe", "action": "ANONYMIZED"}] + } + } + ], + } + return response + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _clean_httpx_response() + return _masked_httpx_response() + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=request_data, + call_type="acompletion", + ) + + assert call_count == 3 + updated_messages = request_data["messages"] + assert updated_messages[0]["content"] == "clean chunk with nothing to mask" + assert updated_messages[1]["content"] == "chunk with PII: [NAME]" + + +@pytest.mark.asyncio +async def test_apply_guardrail_accepted_content_costs_exactly_one_call(): + """Content AWS accepts must cost exactly one ApplyGuardrail call, however far over + the chunk budget it is. Chunking is a recovery path, not something every request + pays for. Regression for: bin-packing eagerly on every request, which split + conversations AWS was happy to take whole and multiplied billed calls and guardrail + latency on traffic that never had a size problem.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + item_text = "x" * (BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2) + messages = [{"role": "user", "content": item_text} for _ in range(3)] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + return _passing_bedrock_httpx_response(f"batch-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 1 + assert result.get("action") == "NONE" + + +@pytest.mark.asyncio +async def test_apply_guardrail_small_content_makes_exactly_one_call(): + """Content that fits entirely within the budget in a single batch must + make exactly one ApplyGuardrail call -- confirms bin-packing does not + introduce an extra probe call for the common (small-request) case.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "short message one"}, + {"role": "user", "content": "short message two"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _passing_bedrock_httpx_response("single-batch") + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_post.assert_awaited_once() + assert result.get("action") == "NONE" + + +@pytest.mark.asyncio +async def test_apply_guardrail_batch_under_budget_still_rejected_falls_back_to_bisection(): + """A batch that fits the budget guess but is still rejected by AWS as too large + (a lower real per-account/region/policy cap) must fall back to bisection for that + batch only, and any other batch from the same request that AWS already accepted + must not be re-sent. + + Three half-budget items pack into two batches once the whole-payload probe is + rejected, so the sequence is probe, batch one (rejected), its two halves, batch + two.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + item_text = "x" * (BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2) + messages = [ + {"role": "user", "content": item_text}, + {"role": "user", "content": item_text}, + {"role": "user", "content": item_text}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count in (1, 2): + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response(f"chunk-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 5 + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-3", "chunk-4", "chunk-5"] + + +def test_split_bedrock_content_single_item_splits_on_whitespace_not_mid_word(): + """A single content item whose raw character midpoint would fall inside a + word must instead split at the nearest whitespace, so neither fragment + ends or begins mid-token. Regression for the Veria AI review finding: a + denied word/PII pattern straddling a raw character-midpoint cut could be + truncated on both fragments and scan clean on each, then reassemble into + the original unmasked text -- a detection bypass.""" + text = ("a" * 20) + " " + ("b" * 30) + raw_midpoint = len(text) // 2 + assert text[raw_midpoint] == "b" + content = [BedrockContentItem(text=BedrockTextContent(text=text))] + + split_content = BedrockGuardrail._split_bedrock_content(content) + assert split_content is not None + first_half, second_half = split_content + + first_text = first_half[0]["text"]["text"] + second_text = second_half[0]["text"]["text"] + + assert first_text + second_text == text + assert first_text == ("a" * 20) + " " + assert second_text == "b" * 30 + + +def test_split_bedrock_content_single_item_with_no_whitespace_falls_back_to_midpoint(): + """A single giant token with no whitespace anywhere has no safe split + point, so the split must fall back to the raw character midpoint rather + than failing or looping.""" + text = "a" * 40 + content = [BedrockContentItem(text=BedrockTextContent(text=text))] + + split_content = BedrockGuardrail._split_bedrock_content(content) + assert split_content is not None + first_half, second_half = split_content + + first_text = first_half[0]["text"]["text"] + second_text = second_half[0]["text"]["text"] + assert first_text + second_text == text + assert len(first_text) == 20 + assert len(second_text) == 20 + + +def test_bin_pack_bedrock_content_packs_minimal_batches_within_budget(): + """Many medium items should pack into the minimal number of in-order + batches that each stay within budget, not one batch per item.""" + items = [BedrockContentItem(text=BedrockTextContent(text="x" * 30)) for _ in range(10)] + + batches = BedrockGuardrail._bin_pack_bedrock_content(items, budget=100) + + assert sum(len(batch) for batch in batches) == 10 + for batch in batches: + combined_len = sum(len(item["text"]["text"]) for item in batch) + assert combined_len <= 100 + assert len(batches) == 4 + + +def test_bin_pack_bedrock_content_oversized_single_item_becomes_its_own_batch(): + """An item whose own text already exceeds the budget must not be + pre-split here -- it becomes its own oversized batch, and only the + reactive bisection fallback (on an AWS rejection) may split it later.""" + small_item = BedrockContentItem(text=BedrockTextContent(text="short")) + oversized_item = BedrockContentItem(text=BedrockTextContent(text="x" * 200)) + items = [small_item, oversized_item, small_item] + + batches = BedrockGuardrail._bin_pack_bedrock_content(items, budget=100) + + assert batches == ((small_item,), (oversized_item,), (small_item,)) + + +def test_bin_pack_bedrock_content_empty_content_makes_exactly_one_empty_batch(): + """Empty content must still pack into exactly one (empty) batch, matching + pre-bin-packing behavior of sending the content list as-is in one call -- + bin-packing must not turn an empty request into zero ApplyGuardrail calls.""" + assert BedrockGuardrail._bin_pack_bedrock_content([], budget=100) == ((),) + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_reported_as_429_bisects_without_burning_retries(): + """AWS reports an oversized ApplyGuardrail request as a 429 ThrottlingException + (not the documented 400 ValidationException) when the guardrail has an active + content-filter policy. That is not a transient throttle -- re-posting the same + oversized content can never succeed -- so it must bisect immediately instead of + consuming the exponential-backoff retry budget first. + + Regression for a bug found against a live guardrail: because the throttle retry + only keyed off status 429, every oversized chunk burned all + _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES attempts (each a billed AWS call, + each preceded by a backoff sleep) before bisection got a chance, at every level + of the recursion.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_throttling_httpx_response() + return _passing_bedrock_httpx_response(f"half-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.asyncio.sleep", + new_callable=AsyncMock, + ) as mock_sleep, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + mock_sleep.assert_not_awaited() + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["half-2", "half-3"] + + +def test_chunk_budget_defaults_to_apply_guardrail_per_second_quota(): + """The default budget must track ApplyGuardrail's default quota of 25 text units + (about 1,000 characters each) per second. Packing to that size and posting + sequentially is what stops chunking from trading a size error for a throttle, so + this default is a deliberate match to AWS behaviour rather than an arbitrary + number.""" + assert BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS == 25_000 + assert BedrockGuardrail(guardrailIdentifier="g", guardrailVersion="DRAFT").chunk_budget_chars == 25_000 + + +@pytest.mark.asyncio +async def test_configured_chunk_budget_changes_how_content_is_packed(): + """An account with raised quotas can set a larger `chunk_budget_chars` and have it + actually drive packing once AWS has rejected a payload, spending fewer + ApplyGuardrail calls for the same content instead of being pinned to the + conservative default. + + Four 20,000-character messages are 80,000 characters total, and every call here is + preceded by the one whole-payload probe AWS rejects. At the 25,000 default only one + message fits per batch, so it is the probe plus four; at 50,000 two fit per batch, + so it is the probe plus two.""" + messages = [{"role": "user", "content": "x" * 20_000} for _ in range(4)] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + async def _calls_made_with_budget(budget: int) -> int: + guardrail = BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + chunk_budget_chars=budget, + ) + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + posted = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal posted + posted += 1 + if posted == 1: + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response("ok") + + mock_post.side_effect = _post_side_effect + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + return mock_post.await_count + + assert await _calls_made_with_budget(BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS) == 5 + assert await _calls_made_with_budget(50_000) == 3 + + +def test_split_index_never_produces_an_empty_fragment(): + """Both fragments must be non-empty for every splittable text, so bisection always + makes progress. + + A text whose only qualifying whitespace is its final character is the dangerous + shape: taking that boundary puts the split at len(text), leaving the first fragment + identical to the input that was just rejected and the second empty. The recursion + would then resubmit the unchanged fragment forever and exhaust the stack instead of + scanning or surfacing Bedrock's error.""" + for text in ("ab ", "xxxx ", ("x" * 40) + " ", " ab", "a b", "ab", " "): + split_at = BedrockGuardrail._nearest_whitespace_split_index(text) + assert 0 < split_at < len(text), f"degenerate split {split_at} for {text!r}" + assert text[:split_at] and text[split_at:], f"empty fragment for {text!r}" + assert text[:split_at] + text[split_at:] == text + + +@pytest.mark.asyncio +async def test_oversized_single_item_with_trailing_space_gives_up_instead_of_recursing(): + """An oversized single item whose only space is trailing must bottom out and + surface Bedrock's error, not recurse forever. + + AWS is modelled the way it really behaves, rejecting every attempt, because the + danger is a fragment identical to the input that was just rejected: AWS would + reject it again, and each retry would split it into the same unchanged fragment. + A split that always shrinks the text terminates and re-raises; one that can return + the whole text raises RecursionError instead. The call-count bound is generous: + halving 41 characters down to unsplittable is a handful of attempts, nowhere near + a stack limit.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + messages = [{"role": "user", "content": ("x" * 40) + " "}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = lambda *_a, **_k: _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as excinfo: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert excinfo.value.status_code == 400 + assert mock_post.await_count < 200 + + class TestBedrockOnlyScanNewMessages: """Bedrock apply_guardrail honors only_scan_new_messages: scans only the per-session diff. @@ -3559,14 +4379,10 @@ class TestBedrockIncrementalFlagInteractions: session = {"litellm_session_id": "sess-flags-mask"} with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 1 mock_api.reset_mock() - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 1, "masking mode must re-scan every turn, exactly once" @pytest.mark.asyncio @@ -3586,9 +4402,7 @@ class TestBedrockIncrementalFlagInteractions: assert mock_api.call_count == 2, "incremental attempt + full-scan fallback" assert result["texts"] == ["MASKED q1"], "masked content must be applied" mock_api.reset_mock() - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 2, "no hashes persisted, so the double scan repeats" @pytest.mark.asyncio @@ -3647,9 +4461,7 @@ async def test_moderation_hook_honors_the_mcp_event_type(mode, call_type, should } with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: - mock_api.return_value = MagicMock( - action="NONE", output=[], outputs=[], assessments=[] - ) + mock_api.return_value = MagicMock(action="NONE", output=[], outputs=[], assessments=[]) await guardrail.async_moderation_hook( data=data, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u"), @@ -3707,3 +4519,146 @@ class TestScanOnlyToolResultsWithLatestRoleFilter: assert result["texts"] == ["TOOL-RESULT"] warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args) assert "scan_only_tool_results" in warning_text + + +@pytest.mark.parametrize("separator", ["\n", "\t", "\r\n", "　"]) +def test_split_bedrock_content_splits_on_any_whitespace_not_just_space(separator): + """Regression: the midpoint split must land on any Unicode whitespace, not only an + ASCII space. + + Matching only " " left the boundary unguarded for exactly the payloads that grow + large enough to need splitting: JSON lines, source code, logs and transcripts are + newline or tab delimited. A deny-listed word sitting at the midpoint of one was cut + in half, scanned clean on both fragments, and reassembled intact, which is the + single-token detection bypass the whitespace split exists to close.""" + text = separator.join(["aaaaaaa"] * 4) + separator + "BADWORDXYZ" + separator + separator.join(["bbbbbbb"] * 4) + + first, second = BedrockGuardrail._split_bedrock_content([BedrockContentItem(text=BedrockTextContent(text=text))]) + + first_text = first[0]["text"]["text"] + second_text = second[0]["text"]["text"] + assert first_text + second_text == text, "split must stay lossless" + assert "BADWORDXYZ" in first_text or "BADWORDXYZ" in second_text, "split severed the token" + + +def test_merge_bedrock_responses_preserves_fields_the_merge_has_no_opinion_on(): + """Regression: merging must not drop AWS response fields it does not itself merge. + + The merged response used to be rebuilt from an empty dict holding only action, + outputs, assessments and usage, so actionReason, guardrailCoverage and anything AWS + adds later vanished from the guardrail_json_response the Admin UI renders, on every + ApplyGuardrail request rather than only chunked ones.""" + chunk = BedrockContentChunkResult( + response={ + "action": "NONE", + "actionReason": "No action.", + "guardrailCoverage": {"textCharacters": {"guarded": 41, "total": 41}}, + "usage": {"contentPolicyUnits": 1}, + }, + content=[BedrockContentItem(text=BedrockTextContent(text="hello"))], + fragment_group_size=1, + ) + + merged = BedrockGuardrail._merge_bedrock_guardrail_responses([chunk]) + + assert merged["actionReason"] == "No action." + assert merged["guardrailCoverage"] == {"textCharacters": {"guarded": 41, "total": 41}} + + +def test_merge_bedrock_usage_sums_counters_not_on_the_known_list(): + """Regression: usage counters were summed from a hardcoded list of six keys, so the + ones AWS also returns (contentPolicyImageUnits, the automatedReasoning pair) were + reported as absent no matter what the chunks actually used.""" + chunks = [ + BedrockContentChunkResult( + response={"action": "NONE", "usage": {"contentPolicyImageUnits": units, "contentPolicyUnits": 1}}, + content=[BedrockContentItem(text=BedrockTextContent(text="x"))], + fragment_group_size=1, + ) + for units in (3, 4) + ] + + usage = BedrockGuardrail._merge_bedrock_guardrail_responses(chunks)["usage"] + + assert usage["contentPolicyImageUnits"] == 7 + assert usage["contentPolicyUnits"] == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_exception_inside_200_logs_failure_and_proceeds(): + """Regression: AWS can report a failure inside an HTTP 200 body via Output.__type, + and that must be logged as guardrail_failed_to_respond rather than success. + + Real AWS does this: an unrecognised operation path on bedrock-runtime answers + HTTP 200 with {"Output": {"__type": "com.amazon.coral.service#UnknownOperationException"}}. + Consolidating telemetry had replaced the derived status with a hardcoded "success", + which reported a failed scan as a clean one. The request itself still proceeds, which + is the behaviour of the code before chunking existed.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + exception_response = MagicMock() + exception_response.status_code = 200 + exception_response.json.return_value = { + "Output": {"__type": "com.amazon.coral.service#UnknownOperationException"}, + "Version": "1.0", + } + exception_response.text = json.dumps(exception_response.json.return_value) + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object(guardrail, "add_standard_logging_guardrail_information_to_request_data") as mock_log, + ): + mock_post.return_value = exception_response + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"model": "bedrock-nova-micro"}, + ) + + assert result is not None, "the request proceeds, as it did before chunking existed" + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_apply_guardrail_failure_logs_a_dict_not_a_bare_string(): + """Regression: the consolidated failure logger must log guardrail_json_response as a + dict, the shape the pre-chunking code and the InvokeGuardrailChecks path both use. + + Consolidating telemetry had changed it to a bare string on the ApplyGuardrail path + only, which breaks any consumer that reads it as a mapping and leaves the two paths + in this file inconsistent.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object(guardrail, "add_standard_logging_guardrail_information_to_request_data") as mock_log, + ): + mock_post.side_effect = _raised_bedrock_error(400, "guardrailIdentifier is not valid") + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_log.assert_called_once() + logged = mock_log.call_args.kwargs["guardrail_json_response"] + assert isinstance(logged, dict), f"expected a dict, got {type(logged).__name__}" + assert "error" in logged diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 00ab39357b4..7a2772ce78c 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1040,6 +1040,64 @@ async def test_apply_guardrail_http_status_error_raises(): assert exc_info.value.status_code == 502 +@pytest.mark.asyncio +async def test_apply_guardrail_404_error_includes_troubleshooting_hint(): + """404 responses include a troubleshooting hint for self-hosted Headroom deployments.""" + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=_make_http_status_error(404, "Not Found"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.detail["status_code"] == 404 + assert exc_info.value.detail["body"] == "Not Found" + assert "hint" in exc_info.value.detail + assert "HEADROOM_COMPRESS_ALLOW_REMOTE=1" in exc_info.value.detail["hint"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_non_404_error_omits_troubleshooting_hint(): + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=_make_http_status_error(500, "headroom internal error"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert exc_info.value.detail["status_code"] == 500 + assert exc_info.value.detail["body"] == "headroom internal error" + assert "hint" not in exc_info.value.detail + + @pytest.mark.asyncio async def test_apply_guardrail_http_status_error_fail_open_forwards_uncompressed(): guardrail = _make_guardrail(unreachable_fallback="fail_open") diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 71e775842e3..8edb56ce25e 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -38,6 +38,41 @@ def test_initialize_presidio_guardrail(): assert result["litellm_params"].mode == "pre_call" +def test_initialize_bedrock_forwards_chunk_budget_chars(): + """Regression: `chunk_budget_chars` set in config.yaml must reach the guardrail. + + The field lives on BedrockGuardrailConfigModel, so LitellmParams parsed it and the + Admin UI rendered it, but initialize_bedrock enumerates its kwargs explicitly and + dropped it. The setting validated and then silently did nothing. Asserting through + initialize_guardrail rather than the constructor is the point: constructing + BedrockGuardrail directly bypasses the only path a user can actually reach. + """ + import litellm + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + test_guardrail = { + "guardrail_name": "test_bedrock_chunk_budget", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.BEDROCK.value, + "mode": "pre_call", + "guardrailIdentifier": "test-guardrail", + "guardrailVersion": "DRAFT", + "chunk_budget_chars": 60_000, + }, + } + + guardrail_handler = InMemoryGuardrailHandler() + guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + initialized = [ + callback + for callback in litellm.callbacks + if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_chunk_budget" + ] + assert initialized, "bedrock guardrail was not registered as a callback" + assert initialized[-1].chunk_budget_chars == 60_000 + + def test_initialize_guardrail_preserves_guardrail_info(): """ Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py index 0074d7062b8..71671966d1a 100644 --- a/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_budget_limiter.py @@ -185,6 +185,35 @@ async def test_team_keys_skip_personal_budget(): mock_get_spend.assert_not_awaited() +@pytest.mark.asyncio +async def test_team_keys_enforce_personal_budget_when_flag_enabled(): + """This hook is the third personal-budget gate alongside common_checks and the + reservation path, so apply_user_budget_to_team_keys has to reach it too or an + opted-in deployment enforces in two places out of three.""" + handler = _PROXY_MaxBudgetLimiter() + user_api_key_dict = _make_user_api_key_auth( + user_max_budget=10.0, + team_id="team-1", + ) + + with patch.dict( + "litellm.proxy.proxy_server.general_settings", + {"apply_user_budget_to_team_keys": True}, + ), patch( + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=999.0), + ): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + + @pytest.mark.asyncio async def test_no_max_budget_passes(): handler = _PROXY_MaxBudgetLimiter() diff --git a/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py b/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py index 0855c201945..61583d11dfa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py +++ b/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py @@ -13,6 +13,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) # --------------------------------------------------------------------------- @@ -227,8 +231,8 @@ async def test_agent_activity_non_admin_no_perms_falls_back_to_owned(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), # no explicit agent permissions + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), # no explicit agent permissions ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -275,8 +279,8 @@ async def test_agent_activity_non_admin_intersects_explicit_agent_ids(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=["agent-permitted"]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-permitted"}))), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -321,8 +325,8 @@ async def test_agent_activity_keyless_caller_does_not_query_created_by_null(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -366,8 +370,8 @@ async def test_agent_activity_non_admin_no_access_returns_empty_page(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 888db031515..3a995e27697 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -292,6 +292,7 @@ class TestAutoRouterBenchmarks: ROW = _SessionAggRow( router_name="live-auto", router_type="complexity", + tier_turns={}, sessions=4, turns=40, unordered_turns=1, @@ -377,6 +378,21 @@ class TestAutoRouterBenchmarks: assert totals.avg_turns_per_session == 10.0 assert totals.spend == 10.0 + def test_tier_names_stay_scoped_to_the_router_type_that_recorded_them(self): + quality = self.ROW.model_copy( + update={"router_name": "quality-auto", "router_type": "quality", "tier_turns": {"2": 7}} + ) + complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}}) + assert complexity.tier_turns == {"medium": 7} + assert quality.tier_turns == {"2": 7} + + def test_summed_totals_carry_no_tier_map_because_names_are_router_scoped(self): + from litellm.proxy.management_endpoints.auto_router_endpoints import _summed_agg_row + + quality = self.ROW.model_copy(update={"router_type": "quality", "tier_turns": {"2": 7}}) + complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}}) + assert _summed_agg_row([complexity, quality]).tier_turns == {} + @pytest.mark.asyncio async def test_non_admin_roles_cannot_read_benchmarks(self): from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks @@ -427,3 +443,26 @@ class TestAutoRouterBenchmarks: assert response.routers_in_scope == 1 assert response.groups[0].router_name == "live-auto" assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "wire_value, expected", [({"simple": 24, "complex": 16}, {"simple": 24, "complex": 16}), ({}, {})] + ) + async def test_the_tier_map_reaches_the_response_as_the_jsonb_column_returns_it( + self, wire_value: dict, expected: dict, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks + + class _DB: + async def query_raw(self, sql: str, *params: object): + return [{**TestAutoRouterBenchmarks.ROW.model_dump(), "tier_turns": wire_value}] + + monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})()) + + response = await get_auto_router_benchmarks( + user_api_key_dict=ADMIN, + start_date="2026-07-01", + end_date="2026-08-01", + ) + assert response.groups[0].tier_turns == expected 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 eedffa1ea5f..1e47010b57c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -932,8 +932,11 @@ async def test_add_team_member_budget_table_success(): ) # Verify the result - assert result == team_info_response assert result.team_member_budget_table == mock_budget_record + assert result == team_info_response.model_copy( + update={"team_member_budget_table": mock_budget_record} + ) + assert team_info_response.team_member_budget_table is None # Verify database call was made correctly mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( @@ -8637,9 +8640,9 @@ class TestBatchResolveAccessGroupResources: with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1"]) - assert sorted(result["ag-1"]["models"]) == ["claude-3", "gpt-4"] - assert result["ag-1"]["mcp_server_ids"] == ["mcp-1"] - assert sorted(result["ag-1"]["agent_ids"]) == ["agent-1", "agent-2"] + assert sorted(result["ag-1"].access_model_names) == ["claude-3", "gpt-4"] + assert result["ag-1"].access_mcp_server_ids == ["mcp-1"] + assert sorted(result["ag-1"].access_agent_ids) == ["agent-1", "agent-2"] @pytest.mark.asyncio async def test_multiple_access_groups(self): @@ -8668,8 +8671,8 @@ class TestBatchResolveAccessGroupResources: with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1", "ag-2"]) - assert result["ag-1"]["models"] == ["gpt-4"] - assert result["ag-2"]["models"] == ["gemini"] + assert result["ag-1"].access_model_names == ["gpt-4"] + assert result["ag-2"].access_model_names == ["gemini"] @pytest.mark.asyncio async def test_missing_access_group_omitted(self): @@ -8735,6 +8738,75 @@ class TestBatchResolveAccessGroupResources: assert "ag-1" in result +class TestResolveTeamAccessGroupResources: + """Tests for the per-team access group resolution on /team/info.""" + + @pytest.mark.asyncio + async def test_populates_flat_lists_and_per_group_details(self): + """access_group_details must attribute each model to the group granting it, + so the UI can show provenance on hover; flat lists stay for back-compat. + Duplicated ids must collapse to one entry (response amplification), and the + input object must stay untouched (resolution returns a copy).""" + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_group_name = "shared-models" + row1.access_model_names = ["gpt-4", "claude-3"] + row1.access_mcp_server_ids = ["mcp-1"] + row1.access_agent_ids = [] + + row2 = MagicMock() + row2.access_group_id = "ag-2" + row2.access_group_name = "extra-models" + row2.access_model_names = ["claude-3", "gemini"] + row2.access_mcp_server_ids = [] + row2.access_agent_ids = ["agent-1"] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( + return_value=[row1, row2] + ) + + team_info = TeamInfoResponseObjectTeamTable( + team_id="team-1", access_group_ids=["ag-1", "ag-2", "ag-1", "ag-missing"] + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resolved = await _resolve_team_access_group_resources(team_info) + + assert team_info.access_group_details is None + assert sorted(resolved.access_group_models or []) == [ + "claude-3", + "gemini", + "gpt-4", + ] + assert resolved.access_group_mcp_server_ids == ["mcp-1"] + assert resolved.access_group_agent_ids == ["agent-1"] + assert [ + (d.access_group_id, d.access_group_name, d.models) + for d in (resolved.access_group_details or []) + ] == [ + ("ag-1", "shared-models", ("gpt-4", "claude-3")), + ("ag-2", "extra-models", ("claude-3", "gemini")), + ] + + @pytest.mark.asyncio + async def test_no_access_groups_leaves_details_unset(self): + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + team_info = TeamInfoResponseObjectTeamTable(team_id="team-1", access_group_ids=[]) + resolved = await _resolve_team_access_group_resources(team_info) + + assert resolved.access_group_details is None + assert resolved.access_group_models is None + + @pytest.mark.asyncio async def test_verify_team_access_denies_unauthorized_user(): """ diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py new file mode 100644 index 00000000000..4a021627c3e --- /dev/null +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -0,0 +1,97 @@ +import os +import sys +from types import MappingProxyType +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy.openai_files_endpoints.common_utils import ( + apply_unified_file_ids, + map_raw_file_ids_to_unified, +) +from litellm.types.utils import LiteLLMBatch + + +def _batch(input_file_id, output_file_id, error_file_id) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-1", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + object="batch", + status="cancelled", + output_file_id=output_file_id, + error_file_id=error_file_id, + ) + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_empty_ids_skips_db(): + prisma_client = MagicMock() + + assert await map_raw_file_ids_to_unified(frozenset(), prisma_client) == {} + + prisma_client.db.litellm_managedfiletable.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_no_prisma_client_returns_empty(): + assert await map_raw_file_ids_to_unified(frozenset({"file-raw-1"}), None) == {} + + +@pytest.mark.asyncio +async def test_map_raw_file_ids_to_unified_bulk_queries_and_filters_to_requested_ids(): + row_a = MagicMock( + unified_file_id="unified-a", + flat_model_file_ids=["file-raw-a", "file-raw-other"], + ) + row_b = MagicMock(unified_file_id="unified-b", flat_model_file_ids=["file-raw-b"]) + prisma_client = MagicMock() + prisma_client.db.litellm_managedfiletable.find_many = AsyncMock(return_value=[row_a, row_b]) + + mapping = await map_raw_file_ids_to_unified( + frozenset({"file-raw-b", "file-raw-a", "file-raw-missing"}), prisma_client + ) + + prisma_client.db.litellm_managedfiletable.find_many.assert_awaited_once_with( + where={"flat_model_file_ids": {"hasSome": ["file-raw-a", "file-raw-b", "file-raw-missing"]}} + ) + assert dict(mapping) == {"file-raw-a": "unified-a", "file-raw-b": "unified-b"} + + +def test_apply_unified_file_ids_swaps_only_mapped_ids(): + batch = _batch(input_file_id="file-raw-in", output_file_id="file-raw-out", error_file_id=None) + + apply_unified_file_ids(batch, MappingProxyType({"file-raw-out": "unified-out"})) + + assert batch.input_file_id == "file-raw-in" + assert batch.output_file_id == "unified-out" + assert batch.error_file_id is None + + +def test_apply_unified_file_ids_swaps_all_three_ids(): + batch = _batch( + input_file_id="file-raw-in", + output_file_id="file-raw-out", + error_file_id="file-raw-err", + ) + + apply_unified_file_ids( + batch, + MappingProxyType( + { + "file-raw-in": "unified-in", + "file-raw-out": "unified-out", + "file-raw-err": "unified-err", + } + ), + ) + + assert (batch.input_file_id, batch.output_file_id, batch.error_file_id) == ( + "unified-in", + "unified-out", + "unified-err", + ) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ac01c6ae1d1..e68e7102fce 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3051,3 +3051,147 @@ def test_list_files_key_allowed_openai_model_still_resolves_team_credentials( mocker, monkeypatch, _team_openai_plus_global_anthropic_router(), ["team-gpt"] ) assert captured_kwargs.get("api_key") == "team-openai-key" + + +@pytest.mark.parametrize( + "http_method, url, patched_litellm_call", + [ + ("get", "/v1/files/file-victim-abc123", "litellm.afile_retrieve"), + ("get", "/v1/files/file-victim-abc123/content", "litellm.afile_content"), + ("delete", "/v1/files/file-victim-abc123", "litellm.afile_delete"), + ], +) +def test_require_managed_files_rejects_raw_provider_file_id( + mocker: MockerFixture, + monkeypatch, + llm_router: Router, + http_method: str, + url: str, + patched_litellm_call: str, +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", True) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_call = mocker.patch(patched_litellm_call, new=mocker.AsyncMock()) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="attacker-user" + ) + + try: + response = getattr(client, http_method)( + url, headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + monkeypatch.setattr("litellm.require_managed_files", False) + + assert response.status_code == 400, response.text + mock_call.assert_not_called() + + +def _unified_managed_file_id() -> str: + import base64 + + from litellm.types.utils import SpecialEnums + + unified_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-3.5-turbo", "file-victim-abc123", "gpt-3.5-turbo-id" + ) + return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + +class _ManagedResourceAccessCheckerStub: + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return True + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return True + + +@pytest.mark.asyncio +async def test_require_managed_files_allows_owned_unified_managed_file_id(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", True) + + await validate_managed_id_requirement( + resource_id=_unified_managed_file_id(), + resource_kind="file", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="owner-user"), + managed_files_obj=_ManagedResourceAccessCheckerStub(), + ) + + +@pytest.mark.asyncio +async def test_managed_file_id_requirement_is_opt_in(monkeypatch): + from litellm.proxy.openai_files_endpoints.common_utils import ( + validate_managed_id_requirement, + ) + + monkeypatch.setattr("litellm.require_managed_files", False) + + await validate_managed_id_requirement( + resource_id="file-victim-abc123", + resource_kind="file", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + managed_files_obj=None, + ) + + +def test_raw_provider_file_id_retrieve_allowed_when_managed_files_not_required( + mocker: MockerFixture, monkeypatch, llm_router: Router +): + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + monkeypatch.setattr("litellm.require_managed_files", False) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + + mock_retrieve = mocker.patch( + "litellm.afile_retrieve", + new=mocker.AsyncMock( + return_value=OpenAIFileObject( + id="file-victim-abc123", + object="file", + bytes=3, + created_at=1234567890, + filename="test.txt", + purpose="user_data", + status="uploaded", + ) + ), + ) + + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="some-user" + ) + + try: + response = client.get( + "/v1/files/file-victim-abc123", headers={"Authorization": "Bearer test-key"} + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert response.status_code == 200, response.text + mock_retrieve.assert_called_once() diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index cf83300ab3b..6ac1e15e7b5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -20,6 +20,7 @@ from __future__ import annotations import asyncio import inspect import json +import logging import os from typing import List, Optional, Union from unittest.mock import AsyncMock, MagicMock, patch @@ -31,6 +32,7 @@ from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps from litellm.proxy.proxy_server import ( + ProxyStartupEvent, _initialize_shared_aiohttp_session, _resolve_pydantic_type, _resolve_typed_dict_type, @@ -728,3 +730,50 @@ def test_otel_global_provider_published_after_callback_init(): "preset logger will not exist yet and a second generic logger will own " "the global provider, orphaning gen-ai spans" ) + + +def test_startup_warns_for_global_budget_without_database(caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=None) + + assert "litellm.max_budget=100.0" in caplog.text + assert "will NOT be enforced" in caplog.text + assert "requests will never be blocked" in caplog.text + + +def test_startup_does_not_warn_for_global_budget_with_database(caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyStartupEvent._warn_budget_without_db(max_budget=100.0, prisma_client=MagicMock()) + + assert "litellm.max_budget" not in caplog.text + + +@pytest.mark.parametrize("max_budget", [0, None]) +def test_startup_does_not_warn_without_global_budget(caplog, max_budget): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + ProxyStartupEvent._warn_budget_without_db(max_budget=max_budget, prisma_client=None) + + assert "litellm.max_budget" not in caplog.text + + +def test_proxy_startup_event_warns_for_global_budget_without_database(): + """Pin the lifespan call that prevents silent DB-less budgets. + + The call must follow Prisma setup so DB-backed deployments do not false-positive. + Direct ``_warn_budget_without_db`` tests cover the warning behavior itself. + """ + wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event) + source = inspect.getsource(wrapped) + budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:") + warn_pos = source.find("_warn_budget_without_db(") + next_startup_section_pos = source.find( + "await ProxyStartupEvent.initialize_scheduled_background_jobs(", + budget_check_pos, + ) + + assert budget_check_pos != -1, "global budget startup block not found" + assert warn_pos != -1, "DB-less budget warning call not found" + assert next_startup_section_pos != -1, "startup section after budget block not found" + assert budget_check_pos < warn_pos < next_startup_section_pos, ( + "DB-less budget warning must run after Prisma setup and the DB-backed budget block" + ) 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 30156849628..91a7e1bc2c2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2638,4 +2638,4 @@ 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() == [] + assert clean_agent_registry.get_agent_list() == () diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index 8b8b7871cce..88dc07e741b 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -582,6 +582,7 @@ def test_public_agent_hub_rewrites_upstream_url_to_proxy(): mock_registry = MagicMock() mock_registry.get_public_agent_list.return_value = [agent] + mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) with ( patch("litellm.public_agent_groups", ["agent-123"]), @@ -631,6 +632,7 @@ def test_public_agent_hub_serializes_http_security_scheme_without_bearer_format( mock_registry = MagicMock() mock_registry.get_public_agent_list.return_value = [agent] + mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) with ( patch("litellm.public_agent_groups", ["agent-123"]), diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1a584423fac..34adb4d2091 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -650,6 +650,47 @@ async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_should_reserve_user_budget_counter_for_team_key_when_flag_enabled(spend_counter_state): + """apply_user_budget_to_team_keys must widen the reservation path too. + + Read-time enforcement alone leaks budget under concurrency, so the opt-in has + to reserve against the personal counter as well or a burst of team-key + requests slips past the owner's max_budget. + """ + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-user-on-team-flagged", + spend=0.0, + user_id="user-on-team-flagged", + team_id="team-no-budget", + ) + team_object = LiteLLM_TeamTable(team_id="team-no-budget", spend=0.0, max_budget=None) + user_object = LiteLLM_UserTable(user_id="user-on-team-flagged", spend=0.0, max_budget=5.0) + + with patch( + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.3, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=team_object, + user_object=user_object, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + apply_user_budget_to_team_keys=True, + ) + + assert counter_cache.in_memory_cache.get_cache(key="spend:user:user-on-team-flagged") == pytest.approx(0.3) + + await release_budget_reservation(reservation) + + @pytest.mark.asyncio async def test_should_seed_org_counter_from_with_budget_cache(spend_counter_state): counter_cache, key_cache = spend_counter_state 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 0e9aac7bf85..6d6fd2e5507 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -20,6 +20,7 @@ from litellm.proxy.litellm_pre_call_utils import ( _get_dynamic_logging_metadata, _get_enforced_params, _get_metadata_variable_name, + _promoted_trace_control_fields, _resolve_credential_from_model_config, _resolve_provider_from_deployment, _update_model_if_key_alias_exists, @@ -5869,4 +5870,183 @@ async def test_key_level_callback_vars_survive_the_strip(): ) assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "key-dd-key", "dd_site": "us5.datadoghq.com"} - assert updated["dd_site"] == "us5.datadoghq.com" \ No newline at end of file + assert updated["dd_site"] == "us5.datadoghq.com" + + +class TestPromotedTraceControlFields: + """LIT-5137: caller metadata trace fields must reach litellm_metadata.""" + + def _make_request(self, path: str) -> MagicMock: + request = MagicMock(spec=Request) + request.url = MagicMock() + request.url.path = path + request.url.__str__.return_value = f"http://localhost{path}" + request.method = "POST" + request.query_params = {} + request.headers = {"Content-Type": "application/json"} + request.client = MagicMock() + request.client.host = "127.0.0.1" + return request + + async def _run(self, path: str, data: dict, headers: dict | None = None) -> dict: + request = self._make_request(path) + if headers is not None: + request.headers = {"Content-Type": "application/json", **headers} + return await add_litellm_data_to_request( + data=data, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + def test_returns_litellm_metadata_for_responses_route(self): + assert _get_metadata_variable_name(self._make_request("/v1/responses")) == "litellm_metadata" + + def test_promotes_trace_prefixed_and_allow_listed_fields(self): + requester_metadata = { + "trace_id": "trace-1", + "trace_name": "name-1", + "trace_user_id": "user-1", + "trace_metadata": {"tenant_id": "tenant-1"}, + "trace_version": "v1", + "trace_release": "r1", + "session_id": "session-1", + "mask_input": True, + "mask_output": True, + } + + promoted = _promoted_trace_control_fields( + requester_metadata=requester_metadata, + litellm_metadata={}, + ) + + assert dict(promoted) == requester_metadata + + def test_does_not_promote_unlisted_trace_prefixed_fields(self): + """trace_public flips a trace to publicly readable, so the allow-list is explicit.""" + promoted = _promoted_trace_control_fields( + requester_metadata={"trace_id": "trace-1", "trace_public": True, "trace_tags": ["a"]}, + litellm_metadata={}, + ) + + assert dict(promoted) == {"trace_id": "trace-1"} + + def test_does_not_promote_non_trace_fields(self): + promoted = _promoted_trace_control_fields( + requester_metadata={ + "trace_id": "trace-1", + "tags": ["free-tier"], + "user_api_key": "forged", + "user_api_key_user_id": "forged-user", + "spend_logs_metadata": {"forged": True}, + "guardrails": ["disabled"], + "debug_langfuse": True, + "session": "not-session-id", + "existing_trace_id": "victim-trace", + "update_trace_keys": ["input", "output"], + }, + litellm_metadata={}, + ) + + assert dict(promoted) == {"trace_id": "trace-1"} + + def test_does_not_promote_trace_mutation_controls(self): + """existing_trace_id + update_trace_keys let a caller overwrite any trace in the project.""" + promoted = _promoted_trace_control_fields( + requester_metadata={ + "trace_id": "trace-1", + "existing_trace_id": "someone-elses-trace", + "update_trace_keys": ["input", "output"], + }, + litellm_metadata={}, + ) + + assert dict(promoted) == {"trace_id": "trace-1"} + + def test_existing_litellm_metadata_value_wins(self): + promoted = _promoted_trace_control_fields( + requester_metadata={"trace_id": "from-body", "session_id": "from-body", "trace_name": "from-body"}, + litellm_metadata={"trace_id": "from-header", "session_id": "from-header"}, + ) + + assert dict(promoted) == {"trace_name": "from-body"} + + def test_empty_requester_metadata_promotes_nothing(self): + assert _promoted_trace_control_fields(requester_metadata={}, litellm_metadata={}) == () + + @pytest.mark.asyncio + async def test_responses_route_end_to_end(self): + caller_metadata = { + "trace_id": "22662678-30c1-41a1-a24b-216d6e5fb83d", + "session_id": "218af06c-28a2-4705-8a0a-5f9970d39326", + "trace_user_id": "user-123", + "trace_metadata": {"tenant_id": "tenant-1"}, + "mask_input": True, + } + + updated = await self._run( + "/v1/responses", + {"model": "gpt-4.1-mini", "input": "say resp", "metadata": copy.deepcopy(caller_metadata)}, + ) + + litellm_metadata = updated["litellm_metadata"] + for key, value in caller_metadata.items(): + assert litellm_metadata[key] == value + assert updated["metadata"] == caller_metadata + + @pytest.mark.asyncio + async def test_messages_route_end_to_end(self): + updated = await self._run( + "/v1/messages", + { + "model": "claude-sonnet-4-5", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"trace_id": "msg-trace-1", "session_id": "msg-session-1"}, + }, + ) + + assert updated["litellm_metadata"]["trace_id"] == "msg-trace-1" + assert updated["litellm_metadata"]["session_id"] == "msg-session-1" + + @pytest.mark.asyncio + async def test_session_id_header_beats_body_metadata(self): + updated = await self._run( + "/v1/responses", + {"model": "gpt-4.1-mini", "input": "say resp", "metadata": {"session_id": "from-body"}}, + headers={"x-litellm-session-id": "from-header-12345678"}, + ) + + assert updated["litellm_metadata"]["session_id"] == "from-header-12345678" + + @pytest.mark.asyncio + async def test_forged_user_api_key_fields_are_not_promoted(self): + updated = await self._run( + "/v1/responses", + { + "model": "gpt-4.1-mini", + "input": "say resp", + "metadata": {"trace_id": "trace-1", "user_api_key_user_id": "forged", "spend_logs_metadata": {"a": 1}}, + }, + ) + + litellm_metadata = updated["litellm_metadata"] + assert litellm_metadata["trace_id"] == "trace-1" + assert litellm_metadata.get("user_api_key_user_id") != "forged" + + @pytest.mark.asyncio + async def test_chat_completions_route_is_untouched(self): + updated = await self._run( + "/v1/chat/completions", + { + "model": "gpt-4.1-mini", + "messages": [{"role": "user", "content": "hello"}], + "metadata": {"trace_id": "trace-1", "session_id": "session-1"}, + }, + ) + + assert "litellm_metadata" not in updated + assert updated["metadata"]["trace_id"] == "trace-1" + assert updated["metadata"]["session_id"] == "session-1" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index efd2ccb3e53..580a58885d9 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7043,6 +7043,39 @@ async def test_update_general_settings_store_model_in_db_false(): assert ps.general_settings["store_model_in_db"] is False +@pytest.mark.asyncio +async def test_update_general_settings_propagates_apply_user_budget_to_team_keys(): + """The Admin UI toggle writes to the DB config, so the flag has to be in the + runtime propagation allowlist. The reverted skip_user_budget_on_team_key was + exposed in /config/list but never propagated, so its toggle did nothing.""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + with patch("litellm.proxy.proxy_server.general_settings", {}): + await proxy_config._update_general_settings(db_general_settings={"apply_user_budget_to_team_keys": "true"}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["apply_user_budget_to_team_keys"] is True + + +@pytest.mark.asyncio +async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins(): + """A DB value must not silently override an explicit YAML setting on reload.""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + proxy_config._yaml_general_settings_keys = {"apply_user_budget_to_team_keys"} + + with patch("litellm.proxy.proxy_server.general_settings", {"apply_user_budget_to_team_keys": True}): + await proxy_config._update_general_settings(db_general_settings={"apply_user_budget_to_team_keys": False}) + + import litellm.proxy.proxy_server as ps + + assert ps.general_settings["apply_user_budget_to_team_keys"] is True + + @pytest.mark.asyncio @pytest.mark.parametrize( "db_value,expected", @@ -9536,6 +9569,38 @@ def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): app.dependency_overrides.clear() +def test_get_config_list_includes_apply_user_budget_to_team_keys(monkeypatch): + """Related to #12905: the opt-in must be discoverable via /config/list so it + renders as a Boolean toggle on the Admin UI General Settings table. This needs + both the ConfigGeneralSettings field and the allowed_args entry.""" + import types + from unittest.mock import AsyncMock, MagicMock + + from fastapi.testclient import TestClient + + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.proxy_server import app + + mock_prisma = MagicMock() + mock_config_table = MagicMock() + mock_config_table.find_first = AsyncMock(return_value=None) + mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + try: + client = TestClient(app) + resp = client.get("/config/list", params={"config_type": "general_settings"}) + assert resp.status_code == 200, resp.text + fields = {item["field_name"]: item for item in resp.json()} + assert "apply_user_budget_to_team_keys" in fields + assert fields["apply_user_budget_to_team_keys"]["field_type"] == "Boolean" + finally: + app.dependency_overrides.clear() + + def test_get_config_list_includes_budget_exceeded_throttle_percentage(monkeypatch): """The throttle fraction is a litellm_settings scalar surfaced on the General Settings table as a Float field so it sits with the other global limits; it diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index abd6220144b..a8e81e92ebd 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1096,6 +1096,7 @@ async def test_prisma_health_check_failure_names_itself_at_operator_visible_leve function and reads as "the check never ran", and reporting it only at debug level hides a database fault behind a flag nobody enables in production.""" import logging + from functools import partial from unittest.mock import AsyncMock from litellm.proxy.utils import PrismaClient @@ -1103,6 +1104,9 @@ async def test_prisma_health_check_failure_names_itself_at_operator_visible_leve client = MagicMock() client.db.query_raw = AsyncMock(side_effect=Exception("connection refused")) client.proxy_logging_obj.failure_handler = AsyncMock() + client._probe_target_wrapper = MagicMock(return_value=client.db) + client._run_health_probe = partial(PrismaClient._run_health_probe, client) + client._report_health_check_failure = AsyncMock() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): with pytest.raises(Exception, match="connection refused"): @@ -1142,6 +1146,7 @@ async def test_prisma_health_check_failure_redacts_database_credentials(caplog): text can carry a full connection string, so the credential has to be gone from the emitted record.""" import logging + from functools import partial from unittest.mock import AsyncMock from litellm.proxy.utils import PrismaClient @@ -1151,6 +1156,9 @@ async def test_prisma_health_check_failure_redacts_database_credentials(caplog): side_effect=Exception("could not connect to postgresql://admin:hunter2@db.internal:5432/litellm") ) client.proxy_logging_obj.failure_handler = AsyncMock() + client._probe_target_wrapper = MagicMock(return_value=client.db) + client._run_health_probe = partial(PrismaClient._run_health_probe, client) + client._report_health_check_failure = AsyncMock() with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): with pytest.raises(Exception): diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 59268e1427b..5fb4392eec6 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -9,6 +9,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ( create_model_info_response, get_available_models_for_user, + hash_token, is_known_model, is_known_vector_store_index, model_dump_with_preserved_fields, @@ -404,3 +405,119 @@ async def test_get_available_models_for_user_error_path_complete_list_raises( general_settings={}, user_model=None, ) + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_resolves_team_access_group_models( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + + team = LiteLLM_TeamTableCachedObj( + team_id="team-1", + models=["no-default-models"], + access_group_ids=["ag-1"], + ) + access_group = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="repro-group", + access_model_names=["model-a", "model-b"], + assigned_team_ids=["team-1"], + ) + + async def _get_team_object(**_kwargs): + return team + + async def _get_access_object(**_kwargs): + return access_group + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["all-team-models"], + team_models=["no-default-models"], + ), + llm_router=_router_with_models(["model-a", "model-b", "model-c"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert sorted(result) == ["model-a", "model-b"] + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_without_access_groups_grants_nothing( + monkeypatch, +): + from litellm.models.team import LiteLLM_TeamTableCachedObj + + async def _get_team_object(**_kwargs): + return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"]) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["all-team-models"], + team_models=["no-default-models"], + ), + llm_router=_router_with_models(["model-a", "model-b"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert result == [] + +@pytest.mark.asyncio +async def test_get_available_models_for_user_resolves_key_access_group_models( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + + async def _get_team_object(**_kwargs): + return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"]) + + async def _get_access_object(**_kwargs): + return LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="key-group", + access_model_names=["model-b"], + assigned_key_ids=[hash_token("sk-test-key")], + ) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["no-default-models"], + team_models=["no-default-models"], + access_group_ids=["ag-1"], + ), + llm_router=_router_with_models(["model-a", "model-b"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert result == ["model-b"] diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py b/tests/test_litellm/proxy/vector_store_files_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py new file mode 100644 index 00000000000..da5dd1934e4 --- /dev/null +++ b/tests/test_litellm/proxy/vector_store_files_endpoints/test_endpoints.py @@ -0,0 +1,146 @@ +""" +require_managed_files enforcement for litellm/proxy/vector_store_files_endpoints/endpoints.py + +Every vector-store file route (create, retrieve, content, update, delete) resolves its +caller-supplied file id through _update_request_data_with_managed_file_id before the +provider call, so the guard lives there once and covers all five. + +A raw or forged managed-looking file id has no ownership row, so without the guard it +is attached to a vector store or read back under shared provider credentials. +""" + +import base64 +import os +import sys +from dataclasses import dataclass +from typing import Literal +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from fastapi import HTTPException + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.vector_store_files_endpoints.endpoints import ( + _update_request_data_with_managed_file_id, +) +from litellm.types.utils import SpecialEnums + +RAW_FILE_ID = "file-victim-abc123" +CALLER = UserAPIKeyAuth(api_key="sk-test", user_id="attacker-user", team_id="team-b") + + +@dataclass(frozen=True) +class ManagedResourceAccessCheckerStub: + file_access: Literal["allow", "deny", "missing"] + + async def can_user_call_unified_file_id( + self, + unified_file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + if self.file_access == "missing": + raise HTTPException(status_code=404, detail=f"File not found: {unified_file_id}") + return self.file_access == "allow" + + async def can_user_call_unified_object_id( + self, + unified_object_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + return False + + +def _unified_file_id() -> str: + unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", "victim-unified-id", "gpt-4o-mini", RAW_FILE_ID, "gpt-4o-mini-id" + ) + return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=") + + +async def _resolve( + file_id: str, + file_access: Literal["allow", "deny", "missing"] = "allow", +): + return await _update_request_data_with_managed_file_id( + data={"vector_store_id": "vs-test", "file_id": file_id}, + file_id=file_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=CALLER, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access=file_access), + llm_router=None, + ) + + +@pytest.mark.asyncio +async def test_raw_file_id_rejected_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _resolve(RAW_FILE_ID) + + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_model_encoded_file_id_rejected_when_managed_files_required(): + """encode_file_id_with_model output is client-forgeable and carries no ownership + row, so it is not a managed file id.""" + from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model + + encoded = encode_file_id_with_model(RAW_FILE_ID, "gpt-4o-mini", id_type="file") + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _resolve(encoded) + + assert exc.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_forged_unified_file_id_rejected_without_ownership_record(): + forged_id = _unified_file_id() + data = {"vector_store_id": "vs-test", "file_id": forged_id} + + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _update_request_data_with_managed_file_id( + data=data, + file_id=forged_id, + request=MagicMock(headers={}, query_params={}), + user_api_key_dict=CALLER, + managed_files_obj=ManagedResourceAccessCheckerStub(file_access="missing"), + llm_router=None, + ) + + assert exc.value.status_code == 404 + assert data["file_id"] == forged_id + + +@pytest.mark.asyncio +async def test_other_teams_unified_file_id_rejected(): + with patch.object(litellm, "require_managed_files", True): + with pytest.raises(HTTPException) as exc: + await _resolve(_unified_file_id(), file_access="deny") + + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_owned_unified_file_id_allowed_when_managed_files_required(): + with patch.object(litellm, "require_managed_files", True): + data, original = await _resolve(_unified_file_id()) + + assert original == _unified_file_id() + assert data["file_id"] == RAW_FILE_ID + + +@pytest.mark.asyncio +async def test_raw_file_id_allowed_when_managed_files_not_required(): + with patch.object(litellm, "require_managed_files", False): + data, original = await _resolve(RAW_FILE_ID) + + assert original is None + assert data["file_id"] == RAW_FILE_ID diff --git a/tests/test_litellm/responses/test_responses_api_bridge_flag.py b/tests/test_litellm/responses/test_responses_api_bridge_flag.py index 463af6562f1..f94c31831bf 100644 --- a/tests/test_litellm/responses/test_responses_api_bridge_flag.py +++ b/tests/test_litellm/responses/test_responses_api_bridge_flag.py @@ -16,6 +16,7 @@ sys.path.insert( import litellm from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import Choices, Message, ModelResponse, Usage class TestUseResponsesApiBridgeFlag: @@ -130,6 +131,44 @@ class TestUseResponsesApiBridgeFlag: mock_bridge_handler.assert_called_once() + @patch("litellm.acompletion") + @patch( + "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" + ) + async def test_allowed_openai_params_forwarded_through_bridge( + self, mock_get_config, mock_acompletion + ): + """allowed_openai_params is a named param of responses(), so it must be + explicitly forwarded to the bridge; otherwise litellm.acompletion raises + UnsupportedParamsError for params the caller explicitly allowed.""" + mock_get_config.return_value = litellm.OpenAIResponsesAPIConfig() + mock_acompletion.return_value = ModelResponse( + id="chatcmpl_123", + model="openai/my-custom-model", + choices=[ + Choices( + index=0, + message=Message(role="assistant", content="Answer"), + finish_reason="stop", + ) + ], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + await litellm.aresponses( + model="openai/my-custom-model", + input="Hello", + use_chat_completions_api=True, + allowed_openai_params=["reasoning_effort"], + reasoning={"effort": "high"}, + litellm_logging_obj=MagicMock(), + ) + + mock_acompletion.assert_called_once() + assert mock_acompletion.call_args.kwargs.get("allowed_openai_params") == [ + "reasoning_effort" + ] + @patch("litellm.responses.file_search.emulated_handler._call_aresponses") @patch( "litellm.responses.main.ProviderConfigManager.get_provider_responses_api_config" diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 8e9e32f5898..356556f3563 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3477,6 +3477,23 @@ class TestSessionAffinity: # Pinned to the first turn's model, not re-classified down to SIMPLE. assert second.model == "o1-preview" + @pytest.mark.asyncio + async def test_a_pinned_turn_reports_the_tier_that_serves_it(self, mock_router_instance, session_affinity_config): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = self._request_kwargs("session-1") + await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + pinned = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert pinned.routing_decision["tier"] == "REASONING" + @pytest.mark.asyncio async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config): mock_router_instance.cache = DualCache() @@ -4362,7 +4379,24 @@ class TestRoutingDecisionContents: assert decision is not None assert decision["cause"] == "default_fallback" assert decision["routed_model"] == response.model - assert "tier" not in decision + assert decision.get("tier") == "MEDIUM" + + @pytest.mark.asyncio + async def test_a_default_model_fallback_claims_no_tier(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "default_model": "gpt-4o"}, + ) + response = await router.async_pre_routing_hook( + model="test-complexity-router", + request_kwargs={}, + messages=[{"role": "system", "content": "be nice"}], + ) + assert response is not None + assert response.routing_decision is not None + assert response.routing_decision["cause"] == "default_fallback" + assert "tier" not in response.routing_decision @pytest.mark.asyncio async def test_session_pin_decision(self, mock_router_instance, basic_config): @@ -5907,11 +5941,21 @@ class TestConversationShapeDiscriminator: ) builds = source.split("self._build_routing_decision(")[1:] assert builds - missing = [ - i - for i, block in enumerate(builds) - if "conversation_continuing=conversation_continuing" not in block.split("),")[0] - ] + missing = [] + for i, block in enumerate(builds): + depth = 0 + end = 0 + for j, char in enumerate(block): + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth < 0: + end = j + break + extracted = block[:end] + if "conversation_continuing=conversation_continuing" not in extracted: + missing.append(i) assert not missing, f"routing decisions {missing} do not carry the conversation shape" 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 98a34de295c..e2348d28701 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -3,6 +3,8 @@ import json import pytest from litellm.router_utils.fallback_event_handlers import ( + AttemptedFallbackTargets, + fallback_attempt_key, get_fallback_model_group, run_async_fallback, ) @@ -142,6 +144,208 @@ async def test_run_async_fallback_skips_original_model_group(): assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 +class RecordingFailRouter: + def __init__(self): + self.attempted_models = [] + + def log_retry(self, kwargs, e): + return kwargs + + async def async_function_with_fallbacks(self, *args, **kwargs): + self.attempted_models.append(kwargs.get("model")) + raise RuntimeError("fallback model also failed") + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_model_group_already_attempted(): + """A fallback graph that loops back on itself must not re-attempt a model group that + already failed for this request. Every group in a cycle fails identically, so + revisiting one multiplies the work and the error output without any chance of + succeeding.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="original failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["already-attempted"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_attempts_a_repeated_target_once(): + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model", "fallback-model", "other-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=5, + fallback_depth=0, + ) + + assert router.attempted_models == ["fallback-model", "other-model"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_forwards_attempted_model_groups_to_nested_call(): + """The nested call is where the next hop of the walk decides what to skip, so the + accumulated set has to reach it, carrying both the group that just failed and the + target being attempted.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"earlier-model"})), + ) + + assert router.received_kwargs["attempted_targets"].keys == frozenset( + {"earlier-model", "primary-model", "fallback-model"} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entry", + [ + {"model": "primary-model", "_target_order": 2}, + {"model": "primary-model", "_excluded_deployment_ids": ["dep-1"]}, + ], +) +async def test_run_async_fallback_still_retargets_the_same_group_via_dict_entry(entry): + """Order-based fallback and weighted intra-group failover both re-target the group that + just failed, selecting a different set of deployments inside it. Those entries are dicts + rather than plain names and must survive a guard that skips already-attempted names.""" + router = RecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[entry], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"primary-model"})), + ) + + assert router.received_kwargs["model"] == "primary-model" + + +@pytest.mark.asyncio +async def test_run_async_fallback_skips_a_repeated_dict_target(): + """A client-side fallback list names its targets with dicts, and that list is re-walked + at every level of the recursion, so an entry that carries no request override has to be + recognised as the same attempt as the bare name.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="original failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "already-attempted"}], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == [] + + +@pytest.mark.asyncio +async def test_run_async_fallback_attempts_a_repeated_dict_target_once(): + router = RecordingFailRouter() + entry = {"model": "fallback-model", "messages": [{"role": "user", "content": "shorter"}]} + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[entry, entry, {"model": "other-model"}], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=5, + fallback_depth=0, + ) + + assert router.attempted_models == ["fallback-model", "other-model"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_keeps_a_request_override_distinct_from_the_bare_name(): + """The documented use of the client-side form is to retry a group with different request + params, so an entry carrying an override must survive even when the bare name of that + same group has already been attempted.""" + router = RecordingFailRouter() + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=[ + {"model": "already-attempted", "messages": [{"role": "user", "content": "shorter"}]} + ], + original_model_group="primary-model", + original_exception=RuntimeError("original failed"), + max_fallbacks=3, + fallback_depth=0, + attempted_targets=AttemptedFallbackTargets(frozenset({"already-attempted"})), + ) + + assert router.attempted_models == ["already-attempted"] + + +@pytest.mark.parametrize( + "target, expected", + [ + ("group-a", "group-a"), + ({"model": "group-a"}, "group-a"), + (None, None), + (["group-a"], None), + ], +) +def test_fallback_attempt_key_identity(target, expected): + """A bare name and a `{"model": name}` entry are the same attempt. A shape with no + usable identity returns None and is never skipped, so an unrecognised entry keeps + today's behaviour rather than being silently dropped.""" + assert fallback_attempt_key(target) == expected + + +def test_fallback_attempt_key_gives_a_param_only_entry_its_own_identity(): + """An entry with no `model` re-targets the group currently being attempted with + different request params, so it is a distinct attempt and still needs an identity.""" + key = fallback_attempt_key({"messages": [{"role": "user", "content": "shorter"}]}) + + assert key is not None + assert key != fallback_attempt_key({"messages": [{"role": "user", "content": "other"}]}) + + +def test_fallback_attempt_key_separates_overrides_from_the_bare_name(): + bare = fallback_attempt_key("group-a") + override = fallback_attempt_key({"model": "group-a", "messages": [{"role": "user", "content": "x"}]}) + other_override = fallback_attempt_key({"model": "group-a", "messages": [{"role": "user", "content": "y"}]}) + order_retarget = fallback_attempt_key({"model": "group-a", "_target_order": 2}) + + assert len({bare, override, other_override, order_retarget}) == 4 + + +def test_fallback_attempt_key_is_stable_across_key_order(): + assert fallback_attempt_key({"model": "group-a", "_target_order": 2}) == fallback_attempt_key( + {"_target_order": 2, "model": "group-a"} + ) + + def test_get_fallback_model_group_does_not_mutate_fallbacks(): """A string fallback must be resolved without mutating the caller's fallbacks list, which is the live router config shared across requests.""" 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 7d453c72652..0d063ad14f5 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 @@ -4,6 +4,7 @@ from unittest.mock import Mock import pytest from litellm import Router +from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.proxy._types import UserAPIKeyAuth from litellm.router_utils.common_utils import ( _deployment_supports_web_search, @@ -11,6 +12,7 @@ from litellm.router_utils.common_utils import ( filter_team_based_models, filter_web_search_deployments, resolve_model_group_alias, + truncate_fallback_error_detail, ) @@ -558,3 +560,27 @@ class TestResolveModelGroupAlias: assert router._get_model_from_alias("group-a") == "group-b" assert router._get_model_from_alias("group-item") == "group-b" assert router._get_model_from_alias("group-b") is None + + +class TestTruncateFallbackErrorDetail: + def test_short_detail_is_returned_unchanged(self): + assert truncate_fallback_error_detail("boom") == "boom" + + def test_detail_at_the_limit_is_returned_unchanged(self): + detail = "x" * ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + assert truncate_fallback_error_detail(detail) == detail + + def test_long_detail_is_bounded_and_reports_what_was_dropped(self): + detail = "x" * (ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS + 500) + + truncated = truncate_fallback_error_detail(detail) + + assert truncated.startswith("x" * ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS) + assert truncated.endswith("... [truncated 500 characters]") + assert len(truncated) < len(detail) + + def test_a_megabyte_of_detail_comes_back_small(self): + """The detail is what a fallback level records about the level below it, so it has + 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 diff --git a/tests/test_litellm/test_budget_ratchet_check.py b/tests/test_litellm/test_budget_ratchet_check.py index 1972c1b6386..22d05f4d00d 100644 --- a/tests/test_litellm/test_budget_ratchet_check.py +++ b/tests/test_litellm/test_budget_ratchet_check.py @@ -68,6 +68,50 @@ def test_new_rule_in_head_is_clean(): assert ratchet.regressions_for("b.json", {}, {"new-rule": _spec_of(5)}) == [] +def test_dropped_rule_that_graduated_to_a_hard_failing_config_is_clean(): + base = {"UP006": _spec_of(0)} + assert ratchet.regressions_for("b.json", base, {}, graduated=("UP006",)) == [] + + +def test_graduation_matches_by_prefix_like_ruff_selectors_do(): + base = {"ANN202": _spec_of(865)} + assert ratchet.regressions_for("b.json", base, {}, graduated=("ANN",)) == [] + + +def test_an_unrelated_graduation_does_not_excuse_a_dropped_rule(): + base = {"C901": _spec_of(3)} + regs = ratchet.regressions_for("b.json", base, {}, graduated=("UP006", "SIM118")) + assert [r.rule for r in regs] == ["C901"] + assert "dropped" in regs[0].detail + + +def test_graduation_never_excuses_a_raised_limit(): + base = {"UP006": _spec_of(0)} + regs = ratchet.regressions_for("b.json", base, {"UP006": _spec_of(7)}, graduated=("UP006",)) + assert [r.rule for r in regs] == ["UP006"] + assert "0 -> 7" in regs[0].detail + + +def test_graduated_selectors_come_from_the_paired_ruff_config(): + selectors = ratchet.graduated_selectors("ruff-strict-budget.json") + assert "UP006" in selectors + assert "ANN" not in selectors + + +def test_budgets_without_a_paired_config_can_never_graduate(): + assert ratchet.graduated_selectors("type-discipline-budget.json") == () + assert ratchet.graduated_selectors("basedpyright-code-budget.json") == () + + +def test_a_selector_the_config_also_ignores_does_not_count_as_graduated(): + lint = {"ignore": ["UP006"], "extend-select": ["UP006", "SIM118"]} + assert ratchet.selectors_hard_failed_by(lint) == ("SIM118",) + + +def test_selectors_hard_failed_by_reads_a_config_with_no_ignore_list(): + assert ratchet.selectors_hard_failed_by({"extend-select": ["UP006"]}) == ("UP006",) + + def test_deleted_budget_file_is_a_regression(): regs = ratchet.regressions_for("b.json", {"LIT006": _spec_of(1)}, None) assert [r.rule for r in regs] == ["*"] diff --git a/tests/test_litellm/test_conftest_isolation.py b/tests/test_litellm/test_conftest_isolation.py index 88889ad7740..15183e68f66 100644 --- a/tests/test_litellm/test_conftest_isolation.py +++ b/tests/test_litellm/test_conftest_isolation.py @@ -1,9 +1,15 @@ import litellm +from litellm import Router +from litellm import router as litellm_router_module from litellm import utils as litellm_utils_module CANARY_MODEL = "conftest-isolation-canary-model" +class _CanaryRouterHolder: + router: Router | None = None + + def test_register_model_ledger_entry_is_scoped_to_this_test(): litellm.register_model({CANARY_MODEL: {"litellm_provider": "openai", "input_cost_per_token": 0.001}}) assert CANARY_MODEL in litellm_utils_module._runtime_registered_model_cost @@ -11,3 +17,20 @@ def test_register_model_ledger_entry_is_scoped_to_this_test(): def test_register_model_ledger_entry_was_rolled_back(): assert CANARY_MODEL not in litellm_utils_module._runtime_registered_model_cost + + +def test_live_router_membership_is_scoped_to_this_test(): + _CanaryRouterHolder.router = Router( + model_list=[ + { + "model_name": "conftest-isolation-canary-router", + "litellm_params": {"model": "openai/conftest-isolation-canary-backend", "api_key": "sk-canary"}, + } + ] + ) + assert _CanaryRouterHolder.router in litellm_router_module._live_routers + + +def test_live_router_membership_was_rolled_back(): + assert _CanaryRouterHolder.router is not None + assert _CanaryRouterHolder.router not in litellm_router_module._live_routers diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py index 274d3c517f6..35d98903226 100644 --- a/tests/test_litellm/test_pre_commit_lint.py +++ b/tests/test_litellm/test_pre_commit_lint.py @@ -129,6 +129,126 @@ def _run(repo: Path, bin_dir: Path, extra_env: dict[str, str]) -> subprocess.Com ) +def _commit_all(repo: Path, message: str) -> None: + subprocess.run(["git", "add", "."], cwd=repo, check=True) + subprocess.run( + ["git", "-c", "user.email=t@t", "-c", "user.name=t", "commit", "-qm", message], + cwd=repo, + check=True, + ) + + +def _set_base_ref(repo: Path) -> None: + subprocess.run( + ["git", "update-ref", "refs/remotes/origin/litellm_internal_staging", "HEAD"], + cwd=repo, + check=True, + ) + + +def test_nothing_staged_scopes_to_working_tree_diff_and_runs_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").write_text("x = 2\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_checks_committed_branch_changes(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").write_text("x = 2\n") + _commit_all(repo, "branch change") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged; scoping to the working tree's diff" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_includes_untracked_files_in_scope(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "brand_new.py").write_text("z = 3\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "litellm/brand_new.py" in proc.stdout + assert "linting Python" in proc.stdout + + +def test_nothing_staged_deletion_only_branch_triggers_checks(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "litellm" / "foo.py").unlink() + _commit_all(repo, "delete module") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing to check" not in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" in proc.stdout + assert "ruff format --check" not in proc.stdout + + +def test_staged_deletion_triggers_checks_without_feeding_missing_files_to_tools(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + subprocess.run(["git", "rm", "-q", "litellm/foo.py"], cwd=repo, check=True) + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing staged" not in proc.stdout + assert "linting Python" in proc.stdout + assert "ruff format --check" not in proc.stdout + + +def test_deleted_dashboard_file_still_triggers_dashboard_lint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + (repo / "ui" / "litellm-dashboard" / "src" / "app.ts").unlink() + _commit_all(repo, "delete dashboard file") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "linting dashboard" in proc.stdout + + +def test_nothing_staged_and_no_changes_is_an_explicit_no_op(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + _set_base_ref(repo) + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "nothing to check" in proc.stdout + assert "linting Python" not in proc.stdout + + +def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 1 + assert "cannot resolve the merge base" in proc.stdout + assert "git fetch origin litellm_internal_staging" in proc.stdout + + +def test_partial_staging_warns_which_checks_were_skipped(tmp_path: Path) -> None: + repo, bin_dir = _sandbox(tmp_path) + _commit_all(repo, "base") + (repo / "notes.md").write_text("hi\n") + subprocess.run(["git", "add", "notes.md"], cwd=repo, check=True) + (repo / "litellm" / "foo.py").write_text("x = 4\n") + proc = _run(repo, bin_dir, {}) + assert proc.returncode == 0, proc.stdout + proc.stderr + assert "SKIPPED Python lint (make lint)" in proc.stdout + assert "litellm/foo.py" in proc.stdout + assert "linting Python" not in proc.stdout + + def test_python_dashboard_and_gen_api_blocks_run_concurrently_with_grouped_output(tmp_path: Path) -> None: repo, bin_dir = _sandbox(tmp_path) barrier_dir = tmp_path / "barrier" @@ -159,8 +279,8 @@ def test_full_output_is_saved_to_a_log_file_in_the_git_dir(tmp_path: Path) -> No assert "linting dashboard" in log assert "API types" in log assert "unstaged/untracked changes" in log - assert f"pre-commit: full log: {log_file}" in proc.stdout - assert "pre-commit: full log:" not in log + assert f"check: full log: {log_file}" in proc.stdout + assert "check: full log:" not in log def test_unwritable_log_warns_and_falls_back_to_running_without_one(tmp_path: Path) -> None: @@ -170,7 +290,7 @@ def test_unwritable_log_warns_and_falls_back_to_running_without_one(tmp_path: Pa assert proc.returncode == 0, proc.stdout + proc.stderr assert "linting Python" in proc.stdout assert "output will not be saved" in proc.stderr - assert "pre-commit: full log:" not in proc.stdout + assert "check: full log:" not in proc.stdout failing = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"}) assert failing.returncode == 1 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 4a3395a7d3f..da2d78edb73 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import logging import os import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -14,6 +15,7 @@ sys.path.insert( import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.integrations.custom_logger import CustomLogger def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -7415,3 +7417,136 @@ class TestAutoRouterMaxInputCharsWiring: router = self._router() assert self._registered_auto_router(router).max_input_chars == DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS + + +class _LogCapture(logging.Handler): + def __init__(self, level): + super().__init__(level=level) + self._level = level + self.messages = [] + + def emit(self, record): + if record.levelno == self._level: + self.messages.append(record.getMessage()) + + +class _FallbackAttemptRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.failed_targets = [] + + async def log_failure_fallback_event(self, original_model_group, kwargs, original_exception): + self.failed_targets.append(kwargs.get("model")) + + +def _cyclic_fallback_router(num_retries=0): + groups = ["group-a", "group-b", "group-c", "group-d"] + return litellm.Router( + model_list=[ + { + "model_name": group, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "mock_response": "litellm.InternalServerError", + }, + } + for group in groups + ], + fallbacks=[ + {"group-a": ["group-b", "group-c"]}, + {"group-b": ["group-a", "group-c"]}, + {"group-c": ["group-d"]}, + {"group-d": ["group-b", "group-a"]}, + ], + num_retries=num_retries, + ) + + +async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwargs): + router_logger = logging.getLogger("LiteLLM Router") + previous_level = router_logger.level + router_logger.setLevel(capture.level) + router_logger.addHandler(capture) + if recorder is not None: + litellm.callbacks.append(recorder) + try: + with pytest.raises(litellm.InternalServerError): + await router.acompletion( + model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs + ) + finally: + router_logger.removeHandler(capture) + router_logger.setLevel(previous_level) + if recorder is not None: + litellm.callbacks.remove(recorder) + + +@pytest.mark.asyncio +async def test_cyclic_fallback_graph_does_not_amplify_one_request(): + """A fallback graph whose entries loop back on each other is easy to build by accident, + and every group in the loop fails identically on a deterministic error, so the walk must + not revisit a group and must not re-emit a growing chained traceback at each level. Left + unbounded, one request blocks the event loop long enough for health probes to fail.""" + recorder = _FallbackAttemptRecorder() + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback(_cyclic_fallback_router(), capture, recorder) + + assert sorted(set(recorder.failed_targets)) == ["group-b", "group-c", "group-d"] + assert len(recorder.failed_targets) == len(set(recorder.failed_targets)) + assert not any("Traceback (most recent call last)" in message for message in capture.messages) + assert sum(len(message) for message in capture.messages) < 5_000 + + +@pytest.mark.asyncio +async def test_retry_breadcrumbs_do_not_carry_the_walk_state(): + """log_retry copies every kwarg into previous_models, which reaches spend logs and + logging callbacks. The set of already-attempted groups is router-internal walk state + with no diagnostic value there, and it is the one entry that is not a plain scalar. + A retry has to be configured for the walk state to reach log_retry at all.""" + router = _cyclic_fallback_router(num_retries=1) + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback(router, capture) + + assert router.previous_models, "no retry breadcrumbs were recorded" + assert any( + "fallback_depth" in breadcrumb for breadcrumb in router.previous_models + ), "no breadcrumb carried router walk state, so this test cannot see the leak" + for breadcrumb in router.previous_models: + assert "attempted_targets" not in breadcrumb + + +@pytest.mark.asyncio +async def test_fallback_traceback_stays_available_at_debug_level(): + """Dropping the stack from the ERROR line is only safe because the fallback path still + emits it once per level at DEBUG, which is what an operator needs to diagnose why every + fallback failed. This pins that remaining debug traceback.""" + capture = _LogCapture(logging.DEBUG) + + await _drive_cyclic_fallback(_cyclic_fallback_router(), capture) + + assert any("Traceback (most recent call last)" in message for message in capture.messages) + + +@pytest.mark.asyncio +async def test_fallback_failure_detail_from_upstream_is_bounded(): + """The detail each level records about the level below it is attacker-influenced, since + it carries whatever the upstream error said. It has to be bounded on its own, so a walk + over several groups cannot compound one large message into the log or into the message + handed back to the caller.""" + huge_message = "z" * 50_000 + capture = _LogCapture(logging.ERROR) + + await _drive_cyclic_fallback( + _cyclic_fallback_router(), + capture, + mock_response=litellm.InternalServerError( + message=huge_message, llm_provider="openai", model="group-a" + ), + ) + + 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 diff --git a/tests/test_litellm/test_ruff_strict_gate.py b/tests/test_litellm/test_ruff_strict_gate.py index abdeb6feecc..206207acb09 100644 --- a/tests/test_litellm/test_ruff_strict_gate.py +++ b/tests/test_litellm/test_ruff_strict_gate.py @@ -1,16 +1,24 @@ import importlib.util +import json +import re +import shutil import subprocess +import sys +import tomllib from pathlib import Path import pytest -_MODULE_PATH = Path(__file__).resolve().parents[2] / "scripts" / "ruff_strict_gate.py" +_REPO_ROOT = Path(__file__).resolve().parents[2] +_MODULE_PATH = _REPO_ROOT / "scripts" / "ruff_strict_gate.py" _spec = importlib.util.spec_from_file_location("ruff_strict_gate", _MODULE_PATH) gate = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(gate) Violation = gate.Violation +_ENABLED_BY_RUFF_DEFAULTS = frozenset({"F401"}) + def rule(name, limit): return {name: {"limit": limit}} @@ -151,3 +159,213 @@ def test_base_point_mid_merge_advances_to_the_merged_in_base_tip(tmp_path): repo, _, base_tip = _branched_repo(tmp_path) _git(repo, "merge", "--no-commit", "--no-ff", "main") assert gate.resolve_base_point("main", cwd=repo) == base_tip + + +def _lint_section(config_name: str) -> dict: + return tomllib.loads((_REPO_ROOT / config_name).read_text())["lint"] + + +def _base_external() -> tuple[str, ...]: + return tuple(_lint_section("ruff.toml")["external"]) + + +def _strict_external() -> tuple[str, ...]: + return tuple(_lint_section("ruff-strict.toml")["external"]) + + +def _strict_selected() -> frozenset: + return frozenset(_lint_section("ruff-strict.toml")["select"]) + + +def _prefix_covered(code: str, prefixes: tuple[str, ...]) -> bool: + return any(code.startswith(prefix) for prefix in prefixes) + + +def _selected_by_the_normal_config() -> frozenset: + return frozenset(_lint_section("ruff.toml")["extend-select"]) | _ENABLED_BY_RUFF_DEFAULTS + + +def _budgeted_rules() -> frozenset: + return frozenset(json.loads((_REPO_ROOT / "ruff-strict-budget.json").read_text())) + + +def _ruff_binary() -> str | None: + beside_interpreter = Path(sys.executable).with_name("ruff") + return str(beside_interpreter) if beside_interpreter.exists() else shutil.which("ruff") + + +_RUFF = _ruff_binary() +_needs_ruff = pytest.mark.skipif(_RUFF is None, reason="ruff is not installed in this environment") + + +def _ruff_output_for_noqa(code: str, *extra_args: str) -> str: + proc = subprocess.run( + [ + _RUFF, + "check", + "--no-cache", + "--stdin-filename", + "litellm/types/_external_probe.py", + *extra_args, + "-", + ], + cwd=_REPO_ROOT, + input=f"def _probe(x: int): # noqa: {code}\n return x\n", + capture_output=True, + text=True, + ) + return proc.stdout + + +def test_every_strict_gate_rule_is_protected_from_base_ruf100(): + unprotected = frozenset( + selector + for selector in _strict_selected() + if not _prefix_covered(selector, _base_external()) + and selector not in _selected_by_the_normal_config() + ) + assert unprotected == frozenset(), ( + f"`ruff check` deletes any `# noqa` naming {sorted(unprotected)} as unused, so suppressing " + "one of those strict-gate rules breaks lint. Cover them in ruff.toml's lint.external or " + "enable them in its lint.extend-select." + ) + + +def test_every_selected_rule_keeps_stale_noqa_detection_somewhere(): + policed_by_strict = frozenset( + selector + for selector in _strict_selected() + if not _prefix_covered(selector, _strict_external()) + ) + policed_by_base = frozenset( + selector + for selector in _selected_by_the_normal_config() + if not _prefix_covered(selector, _base_external()) + ) + shadowed = ( + _strict_selected() | _selected_by_the_normal_config() + ) - policed_by_strict - policed_by_base + assert shadowed == frozenset(), ( + f"no config's RUF100 can ever report a stale `# noqa` for {sorted(shadowed)}: every config " + "that selects each of them also shadows it with an external entry. Narrow the external " + "entry in ruff.toml or ruff-strict.toml." + ) + + +_BASE_OWNED_FAMILY = re.compile(r"E[479]\d+|F\d+|T20\d+") +_BASE_OWNED_SINGLES = frozenset({"PGH004", "RUF008", "RUF009", "RUF100"}) + + +@pytest.fixture(scope="module") +def all_ruff_rule_codes() -> frozenset: + listing = subprocess.run( + [_RUFF, "rule", "--all", "--output-format", "json"], + capture_output=True, + text=True, + ) + assert listing.returncode == 0, listing.stderr + return frozenset( + entry["code"] for entry in json.loads(listing.stdout) if "Removed" not in entry["status"] + ) + + +@_needs_ruff +def test_every_base_owned_rule_is_external_or_selected_in_the_strict_config(all_ruff_rule_codes): + base_owned = frozenset( + code + for code in all_ruff_rule_codes + if _BASE_OWNED_FAMILY.fullmatch(code) or code in _BASE_OWNED_SINGLES + ) + stranded = frozenset( + code + for code in base_owned + if code not in _strict_selected() and not _prefix_covered(code, _strict_external()) + ) + assert stranded == frozenset(), ( + f"the strict gate's RUF100 reads a valid `# noqa` for {sorted(stranded)} as unused, the " + "spurious-breach trap ruff-strict.toml's external override exists to prevent. Cover them " + "there." + ) + double_booked = frozenset( + code + for code in base_owned + if code in _strict_selected() and _prefix_covered(code, _strict_external()) + ) + assert double_booked == frozenset(), ( + f"{sorted(double_booked)} are selected by the strict config yet shadowed by its external " + "list, so their stale suppressions can never be reported. Narrow the external entry in " + "ruff-strict.toml." + ) + + +def test_every_budgeted_rule_is_one_the_gate_actually_measures(): + selectors = tuple(_lint_section("ruff-strict.toml")["select"]) + unmeasured = frozenset(code for code in _budgeted_rules() if not code.startswith(selectors)) + assert unmeasured == frozenset(), ( + f"the gate never counts {sorted(unmeasured)}, so their ceilings are dead config that reads " + "as coverage. Either select them in ruff-strict.toml or drop them from the budget." + ) + + +@_needs_ruff +def test_every_strict_selected_rule_is_budgeted_or_hard_failed_by_the_base_config(all_ruff_rule_codes): + strict_enabled = frozenset( + code + for code in all_ruff_rule_codes + if code.startswith(tuple(_lint_section("ruff-strict.toml")["select"])) + ) + base_hard_failed = tuple(_lint_section("ruff.toml")["extend-select"]) + unpoliced = frozenset( + code + for code in strict_enabled + if code not in _budgeted_rules() + and not code.startswith(base_hard_failed) + and code not in _ENABLED_BY_RUFF_DEFAULTS + ) + assert unpoliced == frozenset(), ( + f"nothing enforces {sorted(unpoliced)}: the gate skips rules missing from the budget, and " + "the base config does not hard-fail them. Re-add a budget ceiling or graduate them into " + "ruff.toml's lint.extend-select." + ) + + +@_needs_ruff +def test_a_noqa_for_a_strict_gate_rule_survives_the_normal_ruff_run(): + assert "RUF100" not in _ruff_output_for_noqa("ANN202") + + +@_needs_ruff +def test_the_external_list_is_what_saves_that_noqa(): + assert "RUF100" in _ruff_output_for_noqa("ANN202", "--config", "lint.external=[]") + + +@_needs_ruff +def test_a_stale_noqa_for_a_locally_enabled_rule_is_still_reported(): + assert "RUF100" in _ruff_output_for_noqa("F401") + + +def _ruff_output_for_source(source: str) -> str: + proc = subprocess.run( + [_RUFF, "check", "--no-cache", "--stdin-filename", "litellm/types/_graduate_probe.py", "-"], + cwd=_REPO_ROOT, + input=source, + capture_output=True, + text=True, + ) + return proc.stdout + + +_DEPRECATED_TYPING_ALIAS = "from typing import List # noqa: UP035\n\n\ndef _probe(x: List[int]) -> None: ...\n" + + +@_needs_ruff +def test_a_graduated_rule_now_fails_the_normal_ruff_run_instead_of_waiting_for_the_gate(): + assert "UP006" in _ruff_output_for_source(_DEPRECATED_TYPING_ALIAS) + + +@_needs_ruff +def test_a_graduated_rule_can_still_be_suppressed_without_tripping_unused_noqa(): + suppressed = _DEPRECATED_TYPING_ALIAS.replace("...\n", "... # noqa: UP006\n") + output = _ruff_output_for_source(suppressed) + assert "UP006" not in output + assert "RUF100" not in output diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py new file mode 100644 index 00000000000..1b66863a82f --- /dev/null +++ b/tests/test_litellm/types/test_router.py @@ -0,0 +1,73 @@ +import pytest + +from litellm.types.router import ( + SPECIAL_MODEL_INFO_PARAMS, + Deployment, + LiteLLM_Params, + ModelInfo, +) +from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams + + +def test_model_info_declares_mirrored_pricing_fields(): + """The pricing keys Deployment mirrors onto model_info must be declared fields, not + extras that only survive because ModelInfo sets extra="allow".""" + for field in SPECIAL_MODEL_INFO_PARAMS: + assert field in ModelInfo.model_fields + + info = ModelInfo(id="x", input_cost_per_token=1e-06) + assert info.__pydantic_extra__ == {} + assert info.input_cost_per_token == 1e-06 + + +def test_special_model_info_params_cannot_drift_from_the_mirror(): + assert SPECIAL_MODEL_INFO_PARAMS == tuple(MirroredPricingParams.model_fields) + assert set(SPECIAL_MODEL_INFO_PARAMS) <= set(CustomPricingLiteLLMParams.model_fields) + assert set(SPECIAL_MODEL_INFO_PARAMS) <= set(LiteLLM_Params.model_fields) + + +def test_custom_pricing_params_keeps_every_field_it_had(): + """The mirrored fields moved to a base class; none of them may go missing from + CustomPricingLiteLLMParams, whose model_fields drive custom-pricing detection.""" + for field in ( + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_character", + "output_cost_per_character", + "cache_read_input_token_cost", + "cache_creation_input_token_cost", + "input_cost_per_second", + "cache_read_input_token_cost_flex", + "input_cost_per_character_above_128k_tokens", + "output_cost_per_audio_token", + ): + assert field in CustomPricingLiteLLMParams.model_fields + + +@pytest.mark.parametrize("field", SPECIAL_MODEL_INFO_PARAMS) +def test_deployment_mirrors_pricing_from_litellm_params_onto_model_info(field): + deployment = Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params(model="gpt-4o", **{field: 3e-06}), + ) + assert getattr(deployment.model_info, field) == 3e-06 + assert deployment.model_info.model_dump(exclude_none=True)[field] == 3e-06 + + +def test_unset_pricing_is_still_absent_from_dumps(): + """/model/info responses and DB writes dump model_info with exclude_none=True, so + declaring the pricing fields must not start emitting ~6 null keys per deployment.""" + dumped = ModelInfo(id="x").model_dump(exclude_none=True) + assert [field for field in SPECIAL_MODEL_INFO_PARAMS if field in dumped] == [] + + +def test_pricing_strings_are_coerced_to_float(): + """Cost values arrive from the DB and the Admin UI as strings; they must land as + floats so cost calculation doesn't multiply a str.""" + info = ModelInfo(id="x", output_cost_per_token="0.000002") + assert info.output_cost_per_token == 2e-06 + + +def test_invalid_pricing_is_rejected(): + with pytest.raises(ValueError): + ModelInfo(id="x", input_cost_per_token="free") diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ab8198304bb..3a670bc7345 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23256 + "limit": 23149 }, "LIT002": { - "limit": 27213 + "limit": 27166 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1091 + "limit": 1086 }, "LIT007": { "limit": 0 @@ -27,9 +27,9 @@ "limit": 0 }, "LIT010": { - "limit": 16783 + "limit": 16760 }, "LIT011": { - "limit": 5602 + "limit": 5598 } } diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index ce7cc9a13db..515a992bc85 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -8613,9 +8613,9 @@ "license": "MIT" }, "node_modules/js-yaml": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.0.tgz", - "integrity": "sha512-1td788aAnnZ5qs7V2QIRl1owjtYpbKt749Y3xauqQgwIIGF/xXWz1wMTEBx5O3LK3lXLVuqXPdPxj2BoFHaW9Q==", + "version": "4.3.1", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.1.tgz", + "integrity": "sha512-CY6crGq313MX8GkwvB7tzgp99vjQxY1++5y10/BKN/GUfHqWaOGQMNZkBvqSzsZKWk/ijwHlWzzkLulsGHhjWQ==", "dev": true, "funding": [ { @@ -10319,9 +10319,9 @@ "license": "MIT" }, "node_modules/nanoid": { - "version": "3.3.16", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.16.tgz", - "integrity": "sha512-bzlKTyNJ7+LdGIIwy8ijFpIqEQIvafahV7eYykJ8Cvh42EdJeODoJ6gUJXpQJvej1BddH8OqTXZNE/KfbWAu8Q==", + "version": "3.3.17", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.17.tgz", + "integrity": "sha512-xQLf0A3HOMlgHq0n247/LRuAOYmB7dXJ/DvAxGvsSBij45XtBSmQycu+F8ODbHwns/XyFZagyL1+J0Offw1E0g==", "funding": [ { "type": "github", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index a9ac92cd023..3953b41a2c2 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -90,7 +90,7 @@ }, "overrides": { "prismjs": "1.30.0", - "js-yaml": "4.3.0", + "js-yaml": "4.3.1", "brace-expansion": "5.0.9", "glob": "13.0.0", "minimatch": "10.2.4", diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index e91b7b73a1e..2f551176091 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -2197,4 +2197,176 @@ describe("EntityUsageExport utils", () => { }); }); }); + + describe("display name resolution from entity metadata", () => { + const entityMetrics = { + spend: 12.25, + api_requests: 40, + successful_requests: 39, + failed_requests: 1, + total_tokens: 900, + prompt_tokens: 500, + completion_tokens: 400, + cache_read_input_tokens: 20, + cache_creation_input_tokens: 10, + }; + + const makeSpendData = (entity: string, metadata?: Record): EntitySpendData => ({ + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + [entity]: { + metrics: entityMetrics, + metadata, + api_key_breakdown: { + key1: { + metrics: entityMetrics, + metadata: { key_alias: "prod-key" }, + }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }); + + it("should export the user email as the entity label and keep the raw user id in the id column", () => { + const result = generateDailyData( + makeSpendData("user-123", { user_email: "ada@example.com", user_alias: "Ada" }), + "User", + ); + + expect(result).toHaveLength(1); + expect(result[0]["User"]).toBe("ada@example.com"); + expect(result[0]["User ID"]).toBe("user-123"); + }); + + it("should fall back to the user alias when the user has no email", () => { + const nullEmail = generateDailyData( + makeSpendData("user-123", { user_email: null, user_alias: "Ada Lovelace" }), + "User", + ); + const missingEmail = generateDailyData(makeSpendData("user-123", { user_alias: "Ada Lovelace" }), "User"); + + expect(nullEmail[0]["User"]).toBe("Ada Lovelace"); + expect(missingEmail[0]["User"]).toBe("Ada Lovelace"); + }); + + it("should fall back to the raw entity key when the entity carries no metadata", () => { + const noMetadata = generateDailyData(makeSpendData("my-tag"), "Tag"); + const emptyMetadata = generateDailyData(makeSpendData("customer-9", {}), "Customer"); + const blankNames = generateDailyData(makeSpendData("user-123", { user_email: null, user_alias: null }), "User"); + + expect(noMetadata[0]["Tag"]).toBe("my-tag"); + expect(emptyMetadata[0]["Customer"]).toBe("customer-9"); + expect(blankNames[0]["User"]).toBe("user-123"); + }); + + it("should prefer the team alias map over any alias in entity metadata", () => { + const result = generateDailyData( + makeSpendData("team-1", { team_alias: "Stale Alias", user_email: "ada@example.com" }), + "Team", + mockTeamAliasMap, + ); + + expect(result[0]["Team"]).toBe("Team One"); + }); + + it("should use the team alias from entity metadata when the alias map has no entry for the team", () => { + const result = generateDailyData( + makeSpendData("team-9", { team_alias: "Team Nine", user_email: "ada@example.com" }), + "Team", + mockTeamAliasMap, + ); + + expect(result[0]["Team"]).toBe("Team Nine"); + }); + + it("should resolve metadata.alias to the user email in getEntityBreakdown", () => { + const withEmail = getEntityBreakdown( + makeSpendData("user-123", { user_email: "ada@example.com", user_alias: "Ada" }), + ); + const withoutEmail = getEntityBreakdown(makeSpendData("user-123", { user_alias: "Ada" })); + + expect(withEmail[0].metadata.alias).toBe("ada@example.com"); + expect(withEmail[0].metadata.id).toBe("user-123"); + expect(withoutEmail[0].metadata.alias).toBe("Ada"); + }); + + it("should resolve the user email on every key row of the keys scope", () => { + const spendData: EntitySpendData = { + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + "user-123": { + metrics: entityMetrics, + metadata: { user_email: "ada@example.com", user_alias: "Ada" }, + api_key_breakdown: { + key1: { metrics: entityMetrics, metadata: { key_alias: "prod-key" } }, + key2: { metrics: entityMetrics, metadata: { key_alias: "dev-key" } }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }; + + const result = generateDailyWithKeysData(spendData, "User"); + + expect(result).toHaveLength(2); + expect(result.map((r) => r["User"])).toEqual(["ada@example.com", "ada@example.com"]); + expect(result.map((r) => r["User ID"])).toEqual(["user-123", "user-123"]); + expect(result.find((r) => r["Key ID"] === "key1")?.["Key Alias"]).toBe("prod-key"); + expect(result.find((r) => r["Key ID"] === "key2")?.["Key Alias"]).toBe("dev-key"); + }); + + it("should resolve each entity's own email in the models scope", () => { + const spendData: EntitySpendData = { + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + "user-a": { + metrics: entityMetrics, + metadata: { user_email: "ada@example.com", user_alias: "Ada" }, + api_key_breakdown: { key1: { metrics: entityMetrics, metadata: {} } }, + }, + "user-b": { + metrics: entityMetrics, + metadata: { user_email: null, user_alias: "Grace" }, + api_key_breakdown: { key2: { metrics: entityMetrics, metadata: {} } }, + }, + }, + models: { + "claude-sonnet-4-5": { + metrics: entityMetrics, + api_key_breakdown: { + key1: { metrics: entityMetrics, metadata: {} }, + key2: { metrics: entityMetrics, metadata: {} }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }; + + const result = generateDailyWithModelsData(spendData, "User"); + + expect(result).toHaveLength(2); + expect(result.every((r) => r.Model === "claude-sonnet-4-5")).toBe(true); + expect(result.find((r) => r["User ID"] === "user-a")?.["User"]).toBe("ada@example.com"); + expect(result.find((r) => r["User ID"] === "user-b")?.["User"]).toBe("Grace"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 8ff389ded5e..53d100040d7 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -3,12 +3,18 @@ import type { DateRangePickerValue } from "@tremor/react"; import Papa from "papaparse"; import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types"; -// Resolve display name for an entity. For teams the teamAliasMap provides -// a human-readable alias; for every other entity type the entity key itself -// (tag name, org id, customer id, …) is already the correct label. -const resolveEntityDisplay = (entity: string, teamAliasMap: Record): { id: string; alias: string } => ({ +const resolveEntityDisplay = ( + entity: string, + teamAliasMap: Record, + entityMetadata?: Record, +): { id: string; alias: string } => ({ id: entity, - alias: teamAliasMap[entity] || entity, + alias: + teamAliasMap[entity] || + entityMetadata?.team_alias || + entityMetadata?.user_email || + entityMetadata?.user_alias || + entity, }); // Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py). @@ -68,7 +74,7 @@ export const getEntityBreakdown = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); if (!entitySpend[entity]) { entitySpend[entity] = { @@ -113,7 +119,7 @@ export const generateDailyData = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); dailyBreakdown.push({ Date: day.date, @@ -164,7 +170,7 @@ export const generateDailyWithKeysData = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap); + const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); const apiKeyBreakdown = data.api_key_breakdown || {}; // Iterate through each API key in the breakdown @@ -241,11 +247,13 @@ export const generateDailyWithModelsData = ( spendData.results.forEach((day) => { const dailyEntityModels: { [key: string]: { [key: string]: any } } = {}; + const dailyEntityMetadata: { [key: string]: Record | undefined } = {}; Object.entries(resolveEntities(day.breakdown)).forEach(([entity, entityData]: [string, any]) => { if (!dailyEntityModels[entity]) { dailyEntityModels[entity] = {}; } + dailyEntityMetadata[entity] = entityData.metadata; Object.entries(day.breakdown.models || {}).forEach(([model, modelData]: [string, any]) => { const entityApiKeys = entityData.api_key_breakdown || {}; @@ -282,7 +290,7 @@ export const generateDailyWithModelsData = ( }); Object.entries(dailyEntityModels).forEach(([entity, models]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, dailyEntityMetadata[entity]); Object.entries(models).forEach(([model, metrics]: [string, any]) => { dailyModelBreakdown.push({ diff --git a/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx b/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx index f6d0f38e2db..6691ef6abdb 100644 --- a/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/TeamSSOSettings.tsx @@ -236,7 +236,7 @@ const TeamSSOSettings: React.FC = ({ accessToken }) => { editContent={ update("budget_duration", v)} + onChange={(v) => update("budget_duration", v ?? null)} style={{ maxWidth: 320 }} /> } diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 2bda0f72cec..e8331294972 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -613,6 +613,34 @@ describe("Teams - access_group_ids in team create", () => { ); }); }); + + it("creates a team with no models selected, sending the no-default-models sentinel instead of an empty list", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Group Only Team" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_alias: "Group Only Team", + models: ["no-default-models"], + }), + ); + }); + }); }); describe("Teams - metadata key-value pairs in team create", () => { diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 95642b93019..42ff8bcc4c2 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -42,6 +42,7 @@ interface TeamProps { import DeleteResourceModal from "./common_components/DeleteResourceModal"; import { teamCreateCall } from "./networking"; +import { normalizeTeamModelSelection } from "./team/teamModelAccess"; import { ModelSelect } from "./ModelSelect/ModelSelect"; const canCreateOrManageTeams = ( @@ -351,7 +352,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser } } - await teamCreateCall(accessToken, formValues); + await teamCreateCall(accessToken, { ...formValues, models: normalizeTeamModelSelection(formValues.models) }); NotificationsManager.success("Team created"); await refreshTeams(); form.resetFields(); @@ -618,17 +619,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser label={ Models{" "} - + } - rules={[ - { - required: true, - message: "Please select at least one model", - }, - ]} name="models" > { expect(labels).toEqual(["Anthropic Family", "OpenAI Family", "Custom Configuration"]); }); - it("never lets a wildcard deployment satisfy a preset", async () => { - const wildcard = [{ model_name: "openai-wild", litellm_params: { model: "openai/*" } }]; + it.each([ + ["a wildcard group", "openai/*"], + ["a plain group over a wildcard underlying model", "openai-wild"], + ])("never lets %s satisfy a preset when the hub lists no expansions", async (_label, modelName) => { + const wildcard = [{ model_name: modelName, litellm_params: { model: "openai/*" } }]; mockFetchAvailableModels.mockResolvedValue(groupsFor(wildcard)); mockFetchAllModelDeployments.mockResolvedValue(wildcard); @@ -650,4 +653,59 @@ describe("AddAutoRouterTab", () => { expect(isOptionDisabled(optionByLabel("OpenAI Family")!)).toBe(true); }); }); + + describe("wildcard-matched presets", () => { + const WILDCARD_DEPLOYMENTS = [{ model_name: "someprovider/*", litellm_params: { model: "someprovider/*" } }]; + + const expandedGroupFor = (model: string): string => `someprovider/${model}`; + + const EXPANDED_HUB_GROUPS: ModelGroup[] = [ + { model_group: "someprovider/*", mode: "chat" }, + ...[...new Set(getAllPresets().flatMap((preset) => [...getRequiredModelsInPreset(preset)]))].map((model) => ({ + model_group: expandedGroupFor(model), + mode: "chat", + })), + ]; + + it("enables a preset whose models exist only as wildcard-expanded groups, labeling the match", async () => { + mockFetchAvailableModels.mockResolvedValue(EXPANDED_HUB_GROUPS); + mockFetchAllModelDeployments.mockResolvedValue(WILDCARD_DEPLOYMENTS); + + renderWithProviders(); + openTemplateDropdown(); + + await waitFor(() => { + expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(false); + }); + expect(optionByLabel("Anthropic Family")!.textContent).toContain("Matches your deployments"); + }); + + it("prefills the expanded group names and submits them", async () => { + const user = userEvent.setup(); + mockFetchAvailableModels.mockResolvedValue(EXPANDED_HUB_GROUPS); + mockFetchAllModelDeployments.mockResolvedValue(WILDCARD_DEPLOYMENTS); + + renderWithProviders(); + openTemplateDropdown(); + await waitFor(() => { + expect(isOptionDisabled(optionByLabel("Anthropic Family")!)).toBe(false); + }); + fireEvent.click(optionByLabel("Anthropic Family")!); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "wildcard-router"); + 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]).toMatchObject({ + complexity_router_config: { + tiers: { + SIMPLE: ANTHROPIC_TIERS.SIMPLE.map(expandedGroupFor), + MEDIUM: ANTHROPIC_TIERS.MEDIUM.map(expandedGroupFor), + COMPLEX: ANTHROPIC_TIERS.COMPLEX.map(expandedGroupFor), + REASONING: ANTHROPIC_TIERS.REASONING.map(expandedGroupFor), + }, + }, + }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx index a70b7602e5b..5ac3b8b2b64 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.test.tsx @@ -1,7 +1,9 @@ -import { act, fireEvent, render, screen } from "@testing-library/react"; -import type { ReactNode } from "react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import type { ReactElement, ReactNode } from "react"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; +import { fetchAvailableModels, fetchAvailableModelsForTeam } from "@/components/llm_calls/fetch_models"; import RouterSettingsAccordion, { RouterSettingsAccordionValue } from "./RouterSettingsAccordion"; vi.mock("../networking", () => ({ @@ -9,11 +11,14 @@ vi.mock("../networking", () => ({ })); vi.mock("@/components/llm_calls/fetch_models", () => ({ - fetchAvailableModels: vi.fn().mockResolvedValue([]), + fetchAvailableModels: vi.fn().mockResolvedValue([{ model_group: "global-model" }]), + fetchAvailableModelsForTeam: vi.fn().mockResolvedValue([{ model_group: "openai/*" }, { model_group: "gpt-5" }]), })); vi.mock("../Settings/RouterSettings/Fallbacks/FallbackSelectionForm", () => ({ - FallbackSelectionForm: () => null, + FallbackSelectionForm: ({ availableModels }: { availableModels: string[] }) => ( +
{availableModels.join(",")}
+ ), })); vi.mock("@tremor/react", () => ({ @@ -39,9 +44,19 @@ vi.mock("../router_settings/RouterSettingsForm", () => ({ ), })); +const renderWithQueryClient = (ui: ReactElement) => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + return render(ui, { + wrapper: ({ children }: { children: ReactNode }) => ( + {children} + ), + }); +}; + describe("RouterSettingsAccordion", () => { beforeEach(() => { - vi.useFakeTimers(); + vi.clearAllMocks(); + vi.useFakeTimers({ shouldAdvanceTime: true }); }); afterEach(() => { @@ -58,7 +73,7 @@ describe("RouterSettingsAccordion", () => { it("debounces propagation and calls onChange once with the last value", async () => { const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); - render(); + renderWithQueryClient(); await flushInitialPropagation(onChange); fireEvent.click(screen.getByText("set-least-busy")); @@ -81,9 +96,51 @@ describe("RouterSettingsAccordion", () => { expect(onChange.mock.calls[0][0].router_settings.routing_strategy).toBe("usage-based-routing"); }); + it("offers the team's own models, including team-scoped BYOK ones, when a teamId is given", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("gpt-5,openai/*"); + }); + expect(fetchAvailableModelsForTeam).toHaveBeenCalledWith("test-token", "team-123"); + expect(fetchAvailableModels).not.toHaveBeenCalled(); + }); + + it("falls back to the proxy-wide model listing when no teamId is given", async () => { + renderWithQueryClient(); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("global-model"); + }); + expect(fetchAvailableModelsForTeam).not.toHaveBeenCalled(); + }); + + it("ignores a stale team's model response that resolves after a newer team was selected", async () => { + const resolvers: ((models: { model_group: string }[]) => void)[] = []; + vi.mocked(fetchAvailableModelsForTeam).mockImplementation( + () => new Promise((resolve) => resolvers.push(resolve)) as Promise<{ model_group: string }[]>, + ); + + const { rerender } = renderWithQueryClient(); + await waitFor(() => expect(resolvers).toHaveLength(1)); + + rerender(); + await waitFor(() => expect(resolvers).toHaveLength(2)); + + await act(async () => { + resolvers[1]([{ model_group: "fast-team-model" }]); + resolvers[0]([{ model_group: "slow-team-model" }]); + }); + + await waitFor(() => { + expect(screen.getByTestId("available-models")).toHaveTextContent("fast-team-model"); + }); + expect(screen.getByTestId("available-models")).not.toHaveTextContent("slow-team-model"); + }); + it("does not call onChange when unmounted mid-wait", async () => { const onChange = vi.fn<(value: RouterSettingsAccordionValue) => void>(); - const { unmount } = render(); + const { unmount } = renderWithQueryClient(); await flushInitialPropagation(onChange); fireEvent.click(screen.getByText("set-least-busy")); diff --git a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx index 08b917e302f..56227abe9ea 100644 --- a/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx +++ b/ui/litellm-dashboard/src/components/common_components/RouterSettingsAccordion.tsx @@ -1,12 +1,13 @@ import React, { useEffect, useState, useImperativeHandle, forwardRef, useRef } from "react"; import { TabPanel, TabPanels, TabGroup, TabList, Tab } from "@tremor/react"; +import { useQuery } from "@tanstack/react-query"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { getRouterSettingsCall } from "../networking"; import RouterSettingsForm, { RouterSettingsFormValue } from "../router_settings/RouterSettingsForm"; import { Fallbacks } from "../Settings/RouterSettings/Fallbacks/AddFallbacks"; import { FallbackSelectionForm } from "../Settings/RouterSettings/Fallbacks/FallbackSelectionForm"; import { FallbackGroup } from "../Settings/RouterSettings/Fallbacks/FallbackGroupConfig"; -import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; +import { fetchAvailableModels, fetchAvailableModelsForTeam, ModelGroup } from "@/components/llm_calls/fetch_models"; export interface RouterSettingsAccordionValue { router_settings: { @@ -30,6 +31,7 @@ interface RouterSettingsAccordionProps { value?: RouterSettingsAccordionValue; onChange?: (value: RouterSettingsAccordionValue) => void; modelData?: any; + teamId?: string | null; } export interface RouterSettingsAccordionRef { @@ -39,7 +41,7 @@ export interface RouterSettingsAccordionRef { const PROPAGATE_WAIT_MS = 100; const RouterSettingsAccordion = forwardRef( - ({ accessToken, value, onChange, modelData }, ref) => { + ({ accessToken, value, onChange, modelData, teamId }, ref) => { const [formValue, setFormValue] = useState({ routerSettings: {}, selectedStrategy: null, @@ -47,7 +49,6 @@ const RouterSettingsAccordion = forwardRef([]); const [fallbackGroups, setFallbackGroups] = useState([]); - const [modelInfo, setModelInfo] = useState([]); const [availableRoutingStrategies, setAvailableRoutingStrategies] = useState([]); const [routerFieldsMetadata, setRouterFieldsMetadata] = useState<{ [key: string]: any }>({}); const [routingStrategyDescriptions, setRoutingStrategyDescriptions] = useState<{ [key: string]: string }>({}); @@ -175,21 +176,11 @@ const RouterSettingsAccordion = forwardRef { - if (!accessToken) { - return; - } - const loadModels = async () => { - try { - const uniqueModels = await fetchAvailableModels(accessToken); - setModelInfo(uniqueModels); - } catch (error) { - console.error("Error fetching model info for fallbacks:", error); - } - }; - loadModels(); - }, [accessToken]); + const { data: modelInfo = [] } = useQuery({ + queryKey: ["fallbackAvailableModels", accessToken, teamId ?? null], + queryFn: () => (teamId ? fetchAvailableModelsForTeam(accessToken, teamId) : fetchAvailableModels(accessToken)), + enabled: Boolean(accessToken), + }); // Helper function to build router_settings from current state const buildRouterSettings = (): RouterSettingsAccordionValue["router_settings"] => { diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index aa74bc60aa1..847a6ca1949 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -5,9 +5,10 @@ const { Option } = Select; interface BudgetDurationDropdownProps { value?: string | null; - onChange?: (value: string) => void; + onChange?: (value: string | undefined) => void; className?: string; style?: React.CSSProperties; + placeholder?: string; } const BudgetDurationDropdown: React.FC = ({ @@ -15,6 +16,7 @@ const BudgetDurationDropdown: React.FC = ({ onChange, className = "", style = {}, + placeholder = "n/a", }) => { return (