diff --git a/.agents/skills/database-access/SKILL.md b/.agents/skills/database-access/SKILL.md index 796dfafe4f6..13fb7b7e286 100644 --- a/.agents/skills/database-access/SKILL.md +++ b/.agents/skills/database-access/SKILL.md @@ -1,6 +1,8 @@ --- name: database-access description: Add database queries, extract existing Prisma calls into repositories, or explicitly port queries to Rust and SQLx without expanding migration scope +metadata: + internal: true --- # Database access diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index f6c06f7d83f..158110cbbab 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2,3 +2,4 @@ /litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri /helm/ @yassin-berriai /terraform/ @yassin-berriai +/.github/workflows/zizmor.yml @yuneng-berri diff --git a/.github/workflows/cost-map-guard.yml b/.github/workflows/cost-map-guard.yml index f9172b9551a..05b22017cbb 100644 --- a/.github/workflows/cost-map-guard.yml +++ b/.github/workflows/cost-map-guard.yml @@ -1,7 +1,7 @@ name: Cost map guard -on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed - pull_request_target: +on: + pull_request: branches: - main - "litellm_**" @@ -19,6 +19,7 @@ jobs: steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: + ref: ${{ github.event.pull_request.base.sha }} persist-credentials: false - name: Fetch the pull request head and its merge base id: revisions diff --git a/.github/workflows/interactions-spec-canary.yml b/.github/workflows/interactions-spec-canary.yml index 7ec3defdcb7..a2e4646dce5 100644 --- a/.github/workflows/interactions-spec-canary.yml +++ b/.github/workflows/interactions-spec-canary.yml @@ -19,7 +19,7 @@ jobs: name: interactions-spec-canary if: github.repository == 'BerriAI/litellm' runs-on: ubuntu-latest - timeout-minutes: 15 + timeout-minutes: 30 permissions: contents: read outputs: @@ -39,8 +39,11 @@ jobs: with: version: "0.10.9" + - name: Set up the Rust bridge + uses: ./.github/actions/rust-bridge + - name: Install dependencies - run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy --no-install-project + run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy - name: Run the compliance tests against Google's live spec id: compliance diff --git a/.github/workflows/spend-logs-guard.yml b/.github/workflows/spend-logs-guard.yml index 354a1b7c947..cc3bda9b880 100644 --- a/.github/workflows/spend-logs-guard.yml +++ b/.github/workflows/spend-logs-guard.yml @@ -1,13 +1,9 @@ name: Spend logs guard -on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's diff is read as data and never executed - pull_request_target: +on: + pull_request: branches: - main - pull_request_review: - types: - - submitted - - dismissed permissions: contents: read @@ -23,7 +19,7 @@ jobs: steps: - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 with: - ref: ${{ github.event.pull_request.base.ref }} + ref: ${{ github.event.pull_request.base.sha }} persist-credentials: false - name: Fetch the pull request head, its merge base and its reviews id: revisions diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index b9f8d7e7dfb..89ca60ad9d1 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -561,6 +561,9 @@ jobs: - name: router_code_coverage run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py + - name: router_surface_check + run: uv run --no-sync python ./tests/code_coverage_tests/router_surface_check.py + - name: test_chat_completion_imports run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py diff --git a/.github/workflows/test-merge-smoke.yml b/.github/workflows/test-merge-smoke.yml index f3f09523491..99dbd1c101b 100644 --- a/.github/workflows/test-merge-smoke.yml +++ b/.github/workflows/test-merge-smoke.yml @@ -132,7 +132,7 @@ jobs: windows-sdk: name: windows-sdk runs-on: windows-2025 - timeout-minutes: 10 + timeout-minutes: 30 env: LITELLM_LOCAL_MODEL_COST_MAP: "True" defaults: @@ -154,8 +154,11 @@ jobs: with: version: "0.10.9" + - name: Set up the Rust bridge + uses: ./.github/actions/rust-bridge + - name: Install dependencies - run: .github/scripts/uv_sync_with_retries.sh --frozen --no-install-workspace --group dev --python 3.11 + run: .github/scripts/uv_sync_with_retries.sh --frozen --group dev --python 3.11 - name: Run the Windows tests run: uv run --no-sync python -m pytest --tb=short -v tests/windows_tests diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index b4b7afb9a09..80d5018fe9a 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -137,6 +137,7 @@ jobs: include: - shard: core-utils test-path: |- + tests/unit/accounting tests/unit/batch_completion tests/unit/decisions tests/unit/extraction @@ -244,6 +245,7 @@ jobs: tests/unit/llms/clarifai tests/unit/llms/claude_code tests/unit/llms/cloudflare + tests/unit/llms/codestral tests/unit/llms/codex tests/unit/llms/cohere tests/unit/llms/cometapi @@ -1306,6 +1308,8 @@ jobs: done openai-sdk-compat: name: OpenAI ${{ matrix.openai-version }} compatibility + needs: rust-bridge + if: ${{ !cancelled() }} permissions: contents: read pull-requests: read @@ -1342,10 +1346,18 @@ jobs: with: version: "0.10.9" + - name: Set up the Rust bridge + if: steps.changes.outputs.decision != 'skip' + timeout-minutes: 5 + uses: ./.github/actions/rust-bridge + with: + artifact: ${{ needs.rust-bridge.outputs.artifact }} + - name: Install dependencies if: steps.changes.outputs.decision != 'skip' run: | .github/scripts/uv_sync_with_retries.sh --frozen --group dev --group proxy-dev --extra proxy --no-install-workspace + cp rust-bridge-dist/litellm/rust_bridge/_native.* litellm/rust_bridge/ uv pip install "openai==${OPENAI_VERSION:?}" --exclude-newer-package openai=2026-10-07T00:00:00Z uv run --no-sync python -c 'import os, openai; assert openai.__version__ == os.environ["OPENAI_VERSION"], openai.__version__' diff --git a/.github/workflows/zizmor.yml b/.github/workflows/zizmor.yml index 73f9efb8df9..eba23e795f3 100644 --- a/.github/workflows/zizmor.yml +++ b/.github/workflows/zizmor.yml @@ -27,10 +27,15 @@ jobs: with: persist-credentials: false + - name: Install zizmor + run: | + curl -fsSL --retry 3 -o "$RUNNER_TEMP/zizmor.tar.gz" \ + https://github.com/zizmorcore/zizmor/releases/download/v1.31.1/zizmor-x86_64-unknown-linux-gnu.tar.gz + echo "d185f726bdaf92df78273217dfdd24ebcd6b0ef48cf3ce776ff5762b6e138c99 $RUNNER_TEMP/zizmor.tar.gz" | sha256sum -c - + tar xzf "$RUNNER_TEMP/zizmor.tar.gz" -C "$RUNNER_TEMP" zizmor + - name: Run zizmor - uses: zizmorcore/zizmor-action@5f14fd08f7cf1cb1609c1e344975f152c7ee938d # v0.5.6 - with: - version: "1.24.1" - min-severity: medium - advanced-security: false - annotations: true + env: + GH_TOKEN: ${{ github.token }} + run: | + "$RUNNER_TEMP/zizmor" --no-ignores --no-config --min-severity=medium --format=github . diff --git a/README.md b/README.md index 66b6acb7a71..a7311882bd0 100644 --- a/README.md +++ b/README.md @@ -427,6 +427,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Watsonx Text (`watsonx_text`)](https://docs.litellm.ai/docs/providers/watsonx) | ✅ | ✅ | ✅ | | | | | | | | | [xAI (`xai`)](https://docs.litellm.ai/docs/providers/xai) | ✅ | ✅ | ✅ | | | | | | | | | [Xinference (`xinference`)](https://docs.litellm.ai/docs/providers/xinference) | | | | ✅ | | | | | | | +| [ZeroGPU (`zerogpu`)](https://docs.litellm.ai/docs/providers/zerogpu) | ✅ | ✅ | ✅ | | | | | | | | [**Read the Docs**](https://docs.litellm.ai/docs/) diff --git a/ci_cd/spend_logs_guard.py b/ci_cd/spend_logs_guard.py index a04999c47c0..5e6d38fba57 100644 --- a/ci_cd/spend_logs_guard.py +++ b/ci_cd/spend_logs_guard.py @@ -2,9 +2,12 @@ needs an approving review from the spend-log owner before it can merge. CODEOWNERS can only match file paths, and every schema change edits the same schema.prisma copies, -so a path rule would put every migration on the owner. This gate reads the diff instead: when no -added line in the watched files mentions the table it passes at once, otherwise it passes only once -the owner's latest review on the pull request is an approval. The workflow re-runs on every review. +so a path rule would put every migration on the owner. This gate reads the diff instead: a schema.prisma +hunk counts when it adds or removes a line inside the `model LiteLLM_SpendLogs { ... }` block, a migration +hunk counts when an added line names the table. Without such a hunk it passes at once, otherwise it passes +only once the owner's latest review on the pull request is an approval that still covers every spend-log line +in the head (approvals survive pushes, so a later commit may not add spend-log lines the owner never saw). +Reviews do not trigger the workflow, so the check is re-run by hand after approving. """ from __future__ import annotations @@ -13,8 +16,10 @@ import argparse import re import subprocess import sys -from collections.abc import Sequence +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -25,7 +30,10 @@ WATCHED: Final = ( "litellm-proxy-extras/litellm_proxy_extras/schema.prisma", ) MIGRATIONS_PREFIX: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/" -TABLE: Final = re.compile(r"(? bool: return path in WATCHED or (path.startswith(MIGRATIONS_PREFIX) and path.endswith("/migration.sql")) -def _added_lines(section: str) -> tuple[str, ...]: - body: Final = section.split("\n@@", 1) - return tuple(line[1:] for line in body[1].splitlines() if line.startswith("+")) if len(body) == 2 else () +def model_block(schema: str) -> range: + """1-based line range of `model LiteLLM_SpendLogs { ... }` including both braces, empty when the schema has none.""" + lines: Final = schema.splitlines() + starts: Final = tuple(number for number, line in enumerate(lines, 1) if MODEL_START.match(line)) + if not starts: + return range(0) + ends: Final = tuple(number for number, line in enumerate(lines, 1) if number > starts[0] and line.startswith("}")) + return range(starts[0], (ends[0] if ends else len(lines)) + 1) -def spend_log_hits(diff: str) -> tuple[str, ...]: - """`path: added line` for every added line in a watched file that names the table.""" +def _hunks(section: str) -> tuple[Hunk, ...]: + headers: Final = tuple(HUNK.finditer(section)) + ends: Final = tuple(header.start() for header in headers[1:]) + (len(section),) + return tuple(_hunk(header, section[header.end() : end]) for header, end in zip(headers, ends)) + + +def _hunk(header: re.Match[str], body: str) -> Hunk: + lines: Final = body.splitlines() + return Hunk( + removed=tuple(enumerate((line[1:] for line in lines if line.startswith("-")), int(header.group(1)))), + added=tuple(enumerate((line[1:] for line in lines if line.startswith("+")), int(header.group(2)))), + ) + + +def spend_log_hits(diff: str, blocks: Mapping[str, Block]) -> tuple[str, ...]: + """`path: +line` or `path: -line` for every changed line that touches the table. + + A schema.prisma line counts when its number falls inside that copy's `blocks` entry (head for added lines, + merge base for removed ones), a migration line when it is added and names the table.""" sections: Final = re.split(r"^diff --git [^\n]*\n", diff, flags=re.MULTILINE) - return tuple(f"{path}: {line.strip()}" for path, line in _watched_added_lines(sections) if TABLE.search(line)) - - -def _watched_added_lines(sections: Sequence[str]) -> tuple[tuple[str, str], ...]: - pairs: Final = tuple((_path_of(section), section) for section in sections) + paths: Final = tuple((_path_of(section), section) for section in sections) return tuple( - (path, line) for path, section in pairs if watched(path) for line in _added_lines(section) - ) # comprehension-ok: path then its lines + hit for path, section in paths if watched(path) for hit in _section_hits(path, _hunks(section), blocks) + ) # comprehension-ok: path then its hits + + +def _section_hits(path: str, hunks: Sequence[Hunk], blocks: Mapping[str, Block]) -> tuple[str, ...]: + if path not in blocks: + return tuple(f"{path}: +{line.strip()}" for hunk in hunks for _, line in hunk.added if TABLE.search(line)) + block: Final = blocks[path] + removed: Final = tuple( + f"{path}: -{line.strip()}" for hunk in hunks for number, line in hunk.removed if number in block.base + ) # comprehension-ok: hunk then its lines + added: Final = tuple( + f"{path}: +{line.strip()}" for hunk in hunks for number, line in hunk.added if number in block.head + ) # comprehension-ok: hunk then its lines + return removed + added def _path_of(section: str) -> str: @@ -70,11 +124,33 @@ def _path_of(section: str) -> str: def approved_by(reviews: Sequence[Review], login: str) -> bool: - """True when `login`'s most recent review on the pull request is an approval.""" - states: Final = tuple(review.state for review in reviews if review.user.login == login) + """True when `login`'s most recent review that changes GitHub's approval state is an approval. + + COMMENTED and PENDING reviews leave an earlier approval in force on GitHub, so they are skipped here too. + """ + own: Final = tuple(review.state for review in reviews if review.user.login == login) + states: Final = tuple(state for state in own if state in DECIDING_STATES) return bool(states) and states[-1] == "APPROVED" +def latest_approval(reviews: Sequence[Review], login: str) -> Review | None: + """`login`'s most recent APPROVED review on the pull request, the one GitHub keeps in force after pushes.""" + own: Final = tuple(review for review in reviews if review.user.login == login and review.state == "APPROVED") + return own[-1] if own else None + + +def approval_covers( + approval: Review, head: str, hits: Sequence[str], hits_at: Callable[[str], Sequence[str] | None] +) -> bool: + """True when the approval is on the head, or the head adds no spend-log line the approved commit did not.""" + if approval.commit_id is None: + return False + if approval.commit_id == head: + return True + reviewed: Final = hits_at(approval.commit_id) + return reviewed is not None and set(hits) <= set(reviewed) + + def added_diff(base: str, head: str) -> str: result: Final = subprocess.run( ("git", "diff", "--unified=0", "--no-color", f"{base}..{head}", "--", *WATCHED, MIGRATIONS_PREFIX), @@ -85,6 +161,27 @@ def added_diff(base: str, head: str) -> str: return result.stdout +def file_at(revision: str, path: str) -> str: + """Contents of `path` at `revision`, empty when the file does not exist there.""" + result: Final = subprocess.run(("git", "show", f"{revision}:{path}"), check=False, capture_output=True, text=True) + return result.stdout if result.returncode == 0 else "" + + +def schema_blocks(base: str, head: str) -> Mapping[str, Block]: + return MappingProxyType( + {path: Block(base=model_block(file_at(base, path)), head=model_block(file_at(head, path))) for path in WATCHED} + ) + + +def fetched(commit: str) -> bool: + result: Final = subprocess.run(("git", "fetch", "--no-tags", "--depth=1", "origin", commit), check=False) + return result.returncode == 0 + + +def reviewed_hits(base: str, commit: str) -> tuple[str, ...] | None: + return spend_log_hits(added_diff(base, commit), schema_blocks(base, commit)) if fetched(commit) else None + + class Arguments(BaseModel): model_config = ConfigDict(frozen=True) base: str @@ -100,18 +197,28 @@ def main(argv: Sequence[str]) -> int: parser.add_argument("--approver", required=True, help="GitHub login whose approval unlocks the merge") parser.add_argument("--reviews", required=True, help="JSON file with the pull request's reviews") args: Final = Arguments.model_validate(vars(parser.parse_args(argv))) - hits: Final = spend_log_hits(added_diff(args.base, args.head)) + hits: Final = spend_log_hits(added_diff(args.base, args.head), schema_blocks(args.base, args.head)) if not hits: - sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; nothing to review.\n") + sys.stdout.write( + "No change inside model LiteLLM_SpendLogs in schema.prisma and no added migration line names it; " + "nothing to review.\n" + ) return 0 - if approved_by(_REVIEWS.validate_json(args.reviews.read_bytes()), args.approver): + reviews: Final = _REVIEWS.validate_json(args.reviews.read_bytes()) + approval: Final = latest_approval(reviews, args.approver) if approved_by(reviews, args.approver) else None + if approval is not None and approval_covers( + approval, args.head, hits, lambda commit: reviewed_hits(args.base, commit) + ): sys.stdout.write(f"LiteLLM_SpendLogs change approved by {args.approver}: {', '.join(hits)}\n") return 0 listed: Final = "".join(f" {hit}\n" for hit in hits) - sys.stdout.write( + reason: Final = ( f"This pull request changes LiteLLM_SpendLogs and needs an approving review from @{args.approver}:\n" - f"{listed}The check re-runs when a review is submitted.\n" + if approval is None + else f"LiteLLM_SpendLogs lines changed after @{args.approver} approved {approval.commit_id}, " + f"a new approving review is needed:\n" ) + sys.stdout.write(f"{reason}{listed}After approving, re-run this job.\n") return 1 diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 65e4d50587f..c190e680347 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -32,6 +32,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.llms.base_llm.guardrail_translation.utils import effective_skip_tool_message_for_guardrail from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails._content_utils import iter_tool_json_fields, map_tool_json_text, walk_user_text from litellm.types.guardrails import GuardrailEventHooks @@ -753,7 +754,12 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): _masked_entity_count.set({}) # Covers multimodal list content + Responses-API input. - inspected = walk_user_text(data, self.redact_text, self.redact_tool_payload) + self._redact_prompt(data) + inspected = walk_user_text( + data, + self.redact_text, + self.redact_tool_payload, + skip_tool_message=effective_skip_tool_message_for_guardrail(self), + ) + self._redact_prompt(data) if inspected == 0: # Image-only, empty-text, and unsupported payloads inspected diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 7a07350e2d9..416444664a3 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -28,7 +28,7 @@ import httpx from fastapi import HTTPException from openai import APIConnectionError, AsyncOpenAI from openai.types.file_deleted import FileDeleted -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, Unpack import litellm @@ -53,7 +53,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_metadata, ) from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend -from litellm.llms.base_llm.files.transformation import BaseFileEndpoints +from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, @@ -309,6 +309,7 @@ def _is_xai_batch_output( class _ManagedFileRow(Protocol): unified_file_id: str file_object: OpenAIFileObject + model_mappings: Mapping[str, str] flat_model_file_ids: Sequence[str] storage_backend: Optional[str] storage_url: Optional[str] @@ -362,6 +363,8 @@ class _ManagedResourceDatabase(Protocol): @property def litellm_managedfiletable(self) -> _ManagedFileTableActions: ... + async def query_raw(self, query: str, *args: object) -> object: ... + @property def litellm_managedobjecttable(self) -> _ManagedObjectTableActions: ... @@ -385,10 +388,83 @@ class _RouterFileCallKwargs(TypedDict, total=False): custom_llm_provider: ReadOnly[str | None] +def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase: + return prisma_client.db + + def _managed_file_table(prisma_client: _ManagedResourcePrismaClient) -> _ManagedFileTableActions: return prisma_client.db.litellm_managedfiletable +class _ListedFileRow(Protocol): + @property + def unified_file_id(self) -> str: ... + + @property + def file_object(self) -> object: ... + + @property + def model_mappings(self) -> Mapping[str, object]: ... + + @property + def storage_backend(self) -> str | None: ... + + +class _ScopedFileRow(BaseModel): + model_config = ConfigDict(frozen=True) + + unified_file_id: str + file_object: JsonValue = None + model_mappings: Mapping[str, JsonValue] + storage_backend: str | None = None + + +_SCOPED_FILE_ROWS: Final = TypeAdapter(tuple[_ScopedFileRow, ...]) + + +def _owner_condition(owner_filter: Mapping[str, object]) -> tuple[str, tuple[str, ...]]: + match owner_filter: + case {"OR": [{"created_by": str() as user_id}, {"team_id": str() as team_id}]}: + return '("created_by" = $1 OR "team_id" = $2)', (user_id, team_id) + case {"created_by": str() as user_id}: + return '"created_by" = $1', (user_id,) + case {"team_id": str() as team_id}: + return '"team_id" = $1', (team_id,) + case _: + return ("FALSE", ()) if owner_filter else ("TRUE", ()) + + +async def _scoped_file_rows( + database: _ManagedResourceDatabase, + owner_filter: Mapping[str, object], + model_mapping_keys: frozenset[str], + cursor_id: str | None, + direction: Literal["asc", "desc"], + take: int, +) -> tuple[_ScopedFileRow, ...]: + owner_condition, owner_params = _owner_condition(owner_filter) + keys_param: Final = len(owner_params) + 1 + comparison: Final = "<" if direction == "desc" else ">" + cursor_condition: Final = ( + f'AND ("created_at", "unified_file_id") {comparison} (SELECT "created_at", "unified_file_id" ' + f'FROM "LiteLLM_ManagedFileTable" WHERE "unified_file_id" = ${keys_param + 1}) ' + if cursor_id + else "" + ) + order: Final = direction.upper() + rows: Final = await database.query_raw( + 'SELECT "unified_file_id", "file_object", "model_mappings", "storage_backend" FROM "LiteLLM_ManagedFileTable" ' + f'WHERE {owner_condition} AND ("model_mappings" ?| ${keys_param}::text[] OR "storage_backend" IS NOT NULL) ' + f'{cursor_condition}ORDER BY "created_at" {order}, "unified_file_id" {order} ' + f"LIMIT ${keys_param + (2 if cursor_id else 1)}", + *owner_params, + sorted(model_mapping_keys), + *((cursor_id,) if cursor_id else ()), + take, + ) + return _SCOPED_FILE_ROWS.validate_python(rows) + + def _iter_provider_file_id_pairs( rows: Sequence[_ManagedFileRow], requested_provider_file_ids: frozenset[str], @@ -2265,14 +2341,23 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): user_api_key_dict: UserAPIKeyAuth, limit: Optional[int] = None, after: Optional[str] = None, + model_scope: Optional[FileListModelScope] = None, + before: Optional[str] = None, **data: Dict, ) -> FileListPage: """List the managed files the caller owns, newest first. Pagination is keyset based on ``unified_file_id`` so a key that owns every file on the proxy still reads one bounded page at a time. - ``purpose`` is applied after parsing, because the managed file table - keeps it inside the ``file_object`` blob instead of a column, and rows + ``after`` pages toward older files and ``before`` toward newer ones, + each page still ordered newest first. + ``model_scope`` is applied in the query, so only rows whose + ``model_mappings`` name one of its deployment ids or model names, and + storage-backed rows, leave the database; a storage-backed row then + stays only when its model name is one the router sends to the listed + provider. ``purpose`` is applied after parsing, because the managed + file table keeps it inside the ``file_object`` blob instead of a + column. Rows whose blob will not parse drop out there too, so a chunk of rows can yield fewer matches than the page holds. Successive chunks are read until the page is full or the caller's rows run out, which keeps @@ -2281,9 +2366,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): scan has to continue past it, the chunk widens to ``FILE_LIST_CONTINUATION_CHUNK_SIZE``, so the walk costs one query per that many rows instead of one per page. That bound is per query, not - per request: the work is still linear in the rows the caller owns, and - a filter matching nothing reads every one of them, with no index - covering either the owner filter or the sort. + per request: a ``purpose`` filter matching nothing still reads every + row that passes the query. """ validate_file_list_limit(limit) validate_file_list_purpose(purpose) @@ -2292,36 +2376,60 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if owner_filter is None: return FileListPage.model_validate(build_list_page([])) - if after: - cursor_row = await _managed_file_table(self.prisma_client).find_first( - where={**owner_filter, "unified_file_id": after} + if after and before: + raise ProxyException( + message="Pass either 'after' (or 'after_id') or 'before_id' when listing files, not both.", + type="invalid_request_error", + param="before_id", + code=400, + openai_code="invalid_value", + ) + cursor_param: Final = "before_id" if before else "after" + cursor: Final = before or after + database: Final = _managed_database(self.prisma_client) + if cursor: + cursor_row = await database.litellm_managedfiletable.find_first( + where={**owner_filter, "unified_file_id": cursor} ) if cursor_row is None: raise ProxyException( - message=f"Invalid 'after' cursor: no file found with id '{after}'.", + message=f"Invalid '{cursor_param}' cursor: no file found with id '{cursor}'.", type="invalid_request_error", - param="after", + param=cursor_param, code=400, openai_code="invalid_value", ) + if model_scope is not None and not model_scope.model_mapping_keys: + return FileListPage.model_validate(build_list_page([])) + page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT) matches: Final[List[OpenAIFileObject]] = [] - cursor_id = after + cursor_id = cursor chunk_size = page_size + 1 + direction: Final[Literal["asc", "desc"]] = "asc" if before else "desc" while len(matches) <= page_size: cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {} - chunk = await _managed_file_table(self.prisma_client).find_many( - where=owner_filter, - take=chunk_size, - order=[{"created_at": "desc"}, {"unified_file_id": "desc"}], - **cursor_args, + chunk: Sequence[_ListedFileRow] = ( + await _scoped_file_rows( + database, owner_filter, model_scope.model_mapping_keys, cursor_id, direction, chunk_size + ) + if model_scope is not None + else await database.litellm_managedfiletable.find_many( + where=owner_filter, + take=chunk_size, + order=[{"created_at": direction}, {"unified_file_id": direction}], + **cursor_args, + ) ) matches.extend( _public_file_object(parsed_file_object, row.unified_file_id) for row in chunk - if (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None + if ( + model_scope is None or model_scope.holds(row.model_mappings, stored=row.storage_backend is not None) + ) + and (parsed_file_object := _parse_managed_file_object(row.file_object, row.unified_file_id)) is not None and (purpose is None or parsed_file_object.purpose == purpose) ) if len(chunk) < chunk_size: @@ -2329,7 +2437,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): cursor_id = chunk[-1].unified_file_id chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE) - return FileListPage.model_validate(build_list_page(matches[:page_size], has_more=len(matches) > page_size)) + page: Final = matches[:page_size] + return FileListPage.model_validate( + build_list_page(page[::-1] if before else page, has_more=len(matches) > page_size) + ) def _is_batch_polling_enabled(self) -> bool: """ diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 69068cfc8a5..6dc8bab45ce 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -153,7 +153,7 @@ RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \ find /app/.venv -type d -path "*/tornado/test" -delete && \ chmod -R a+rX /opt/prisma && \ python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \ - python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()" + python -c "import litellm, litellm.rust_bridge._native; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__" USER nonroot diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py index 7b01e64dd73..e22639daaef 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_data_checks.py @@ -343,11 +343,16 @@ def btree_max_entry_bytes(block_size: int) -> int: def _value_bytes(column: ColumnType, identifier: "sql.Identifier") -> "sql.Composed": + """The bytes a value takes in the index entry: its length, or its stored size when + Postgres compressed it, since `index_form_tuple` keeps the compressed bytes.""" from psycopg import sql - if column.type_name in _OCTET_LENGTH_TYPES: - return sql.SQL("octet_length({})").format(identifier) - return sql.SQL("octet_length({}::text)").format(identifier) + length: Final = ( + sql.SQL("octet_length({})").format(identifier) + if column.type_name in _OCTET_LENGTH_TYPES + else sql.SQL("octet_length({}::text)").format(identifier) + ) + return sql.SQL("least({}, pg_column_size({}))").format(length, identifier) def _aligned(offset: "sql.Composable", align: str) -> "sql.Composed": diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261021000000_managed_file_created_by_created_at_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261021000000_managed_file_created_by_created_at_index/migration.sql new file mode 100644 index 00000000000..8bbcb67bddf --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261021000000_managed_file_created_by_created_at_index/migration.sql @@ -0,0 +1,2 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_ManagedFileTable_created_by_created_at_idx" +ON "LiteLLM_ManagedFileTable" ("created_by", "created_at" DESC); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261027000000_add_accounting_ledger_tables/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261027000000_add_accounting_ledger_tables/migration.sql new file mode 100644 index 00000000000..c0ad8692d36 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261027000000_add_accounting_ledger_tables/migration.sql @@ -0,0 +1,47 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE TYPE "AccountingEventKind" AS ENUM ('inbound_usage', 'outbound_usage', 'cost'); + +CREATE TYPE "AccountingDeliveryState" AS ENUM ('pending', 'applied', 'parked'); + +CREATE TYPE "AccountingDeliveryTarget" AS ENUM ('LiteLLM_VerificationToken', 'LiteLLM_UserTable', 'LiteLLM_TeamTable', 'LiteLLM_TeamMembership', 'LiteLLM_OrganizationTable', 'LiteLLM_OrganizationMembership', 'LiteLLM_EndUserTable', 'LiteLLM_AgentsTable', 'LiteLLM_TagTable', 'LiteLLM_ProjectTable', 'LiteLLM_ModelAccessGroupBudgetTable', 'LiteLLM_BudgetTable', 'LiteLLM_BudgetWindowSpend', 'LiteLLM_DailyUserSpend', 'LiteLLM_DailyTeamSpend', 'LiteLLM_DailyOrganizationSpend', 'LiteLLM_DailyEndUserSpend', 'LiteLLM_DailyAgentSpend', 'LiteLLM_DailyTagSpend', 'LiteLLM_DailyModelUsage', 'LiteLLM_DailyToolSpend', 'LiteLLM_AutoRouterDailySpend'); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingEvent" ( + "event_id" TEXT NOT NULL, + "kind" "AccountingEventKind" NOT NULL, + "source_event_id" TEXT, + "schema_version" INTEGER NOT NULL, + "payload_hash" TEXT NOT NULL, + "occurred_at" TIMESTAMP(3) NOT NULL, + "recorded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "payload" JSONB NOT NULL, + + CONSTRAINT "LiteLLM_AccountingEvent_schema_version_check" CHECK ("schema_version" > 0), + CONSTRAINT "LiteLLM_AccountingEvent_pkey" PRIMARY KEY ("event_id") +); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_AccountingDelivery" ( + "event_id" TEXT NOT NULL, + "target" "AccountingDeliveryTarget" NOT NULL, + "version" INTEGER NOT NULL, + "state" "AccountingDeliveryState" NOT NULL DEFAULT 'pending', + "attempts" INTEGER NOT NULL DEFAULT 0, + "next_attempt_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "applied_at" TIMESTAMP(3), + "last_error" TEXT, + + CONSTRAINT "LiteLLM_AccountingDelivery_version_check" CHECK ("version" > 0), + CONSTRAINT "LiteLLM_AccountingDelivery_attempts_check" CHECK ("attempts" >= 0), + CONSTRAINT "LiteLLM_AccountingDelivery_pkey" PRIMARY KEY ("event_id","target","version") +); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_recorded_at_idx" ON "LiteLLM_AccountingEvent"("recorded_at"); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingEvent_kind_source_event_id_idx" ON "LiteLLM_AccountingEvent"("kind", "source_event_id"); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_AccountingDelivery_target_state_next_attempt_at_idx" ON "LiteLLM_AccountingDelivery"("target", "state", "next_attempt_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 83a0e6f9e05..2965174114b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) @@index([flat_model_file_ids], type: Gin) @@index([team_id, created_at(sort: Desc)]) + @@index([created_by, created_at(sort: Desc)]) } model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the @@ -2060,3 +2061,73 @@ model LiteLLM_TelemetryReport { @@index([window_end]) @@index([created_at]) } + +enum AccountingEventKind { + inbound_usage + outbound_usage + cost +} + +enum AccountingDeliveryState { + pending + applied + parked +} + +enum AccountingDeliveryTarget { + LiteLLM_VerificationToken + LiteLLM_UserTable + LiteLLM_TeamTable + LiteLLM_TeamMembership + LiteLLM_OrganizationTable + LiteLLM_OrganizationMembership + LiteLLM_EndUserTable + LiteLLM_AgentsTable + LiteLLM_TagTable + LiteLLM_ProjectTable + LiteLLM_ModelAccessGroupBudgetTable + LiteLLM_BudgetTable + LiteLLM_BudgetWindowSpend + LiteLLM_DailyUserSpend + LiteLLM_DailyTeamSpend + LiteLLM_DailyOrganizationSpend + LiteLLM_DailyEndUserSpend + LiteLLM_DailyAgentSpend + LiteLLM_DailyTagSpend + LiteLLM_DailyModelUsage + LiteLLM_DailyToolSpend + LiteLLM_AutoRouterDailySpend +} + +// Durable accounting ledger. One row per usage or cost event; payload holds the +// private event envelope versioned by schema_version. Usage rows retain request-scoped +// pricing context; cost rows include resolved rates and the finalized pricing snapshot. +model LiteLLM_AccountingEvent { + event_id String @id + kind AccountingEventKind + source_event_id String? + schema_version Int + payload_hash String + occurred_at DateTime + recorded_at DateTime @default(now()) + payload Json + + @@index([recorded_at]) + @@index([kind, source_event_id]) +} + +// One row per (event, target table, version): pending work for the projector and, +// once applied in the same transaction as the target write, the receipt that stops a retry. +model LiteLLM_AccountingDelivery { + event_id String + target AccountingDeliveryTarget + version Int + state AccountingDeliveryState @default(pending) + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + + @@id([event_id, target, version]) + @@index([target, state, next_attempt_at]) +} diff --git a/litellm-rust/.agents/skills/rust-error/SKILL.md b/litellm-rust/.agents/skills/rust-error/SKILL.md index fd5d6222287..74da1ecbcf4 100644 --- a/litellm-rust/.agents/skills/rust-error/SKILL.md +++ b/litellm-rust/.agents/skills/rust-error/SKILL.md @@ -1,6 +1,8 @@ --- name: rust-error description: Define, wrap, convert, classify, report or test Rust error types in litellm-rust with thiserror. Use when adding an error type or variant, mapping one crate's error into another, turning an error into a gateway response or Python exception, deciding retry or fault from an error, or replacing call sites that build errors from strings +metadata: + internal: true --- # Rust errors diff --git a/litellm-rust/.agents/skills/rust-gateway/SKILL.md b/litellm-rust/.agents/skills/rust-gateway/SKILL.md index 66b93ee128d..4fee0b54a68 100644 --- a/litellm-rust/.agents/skills/rust-gateway/SKILL.md +++ b/litellm-rust/.agents/skills/rust-gateway/SKILL.md @@ -1,6 +1,8 @@ --- name: rust-gateway description: Add or change an HTTP route, extractor, response envelope, stream or router mount in the litellm-rust gateway crates (crates/gateway and crates/gateway-*). Use when adding an endpoint, parsing a request body, choosing an error envelope for an API surface, relaying SSE, or writing a router-level test +metadata: + internal: true --- # Rust gateway diff --git a/litellm-rust/.agents/skills/rust-string-enums/SKILL.md b/litellm-rust/.agents/skills/rust-string-enums/SKILL.md index 5fc067790a0..352f6427dd9 100644 --- a/litellm-rust/.agents/skills/rust-string-enums/SKILL.md +++ b/litellm-rust/.agents/skills/rust-string-enums/SKILL.md @@ -1,6 +1,8 @@ --- name: rust-string-enums description: Define or refactor Rust string-valued enums and their Serde adapters in litellm-rust, using Strum and serde_with while preserving parsing, wire values, and schemas +metadata: + internal: true --- # Rust string enums diff --git a/litellm-rust/.agents/skills/rust-tracing/SKILL.md b/litellm-rust/.agents/skills/rust-tracing/SKILL.md index c1c7f43c456..a8da3f3bb14 100644 --- a/litellm-rust/.agents/skills/rust-tracing/SKILL.md +++ b/litellm-rust/.agents/skills/rust-tracing/SKILL.md @@ -1,6 +1,8 @@ --- name: rust-tracing description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery +metadata: + internal: true --- # Rust tracing diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 9c45c5c4b1c..10e590b5006 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4393,6 +4393,10 @@ dependencies = [ "rstest", ] +[[package]] +name = "litellm-router-core" +version = "0.1.0" + [[package]] name = "litellm-router-types" version = "0.1.0" diff --git a/litellm-rust/crates/core-utils/src/core_helpers.rs b/litellm-rust/crates/core-utils/src/core_helpers.rs index ada1c3ceb1a..58ee88d301c 100644 --- a/litellm-rust/crates/core-utils/src/core_helpers.rs +++ b/litellm-rust/crates/core-utils/src/core_helpers.rs @@ -13,6 +13,7 @@ const FINISH_REASONS: &[(&str, &str)] = &[ ("max_tokens", "length"), ("refusal", "content_filter"), ("compaction", "length"), + ("model_context_window_exceeded", "length"), ("guardrail_intervened", "content_filter"), ("content_filtered", "content_filter"), ("content_filter", "content_filter"), @@ -74,6 +75,7 @@ mod tests { #[case::end_turn("end_turn", "stop")] #[case::stop_sequence("stop_sequence", "stop")] #[case::max_tokens("max_tokens", "length")] + #[case::model_context_window_exceeded("model_context_window_exceeded", "length")] #[case::refusal("refusal", "content_filter")] #[case::guardrail_intervened("guardrail_intervened", "content_filter")] #[case::content_filtered("content_filtered", "content_filter")] diff --git a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs index 6517b093bbd..615d530af19 100644 --- a/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs +++ b/litellm-rust/crates/core-utils/src/exception_mapping_utils/original.rs @@ -113,6 +113,7 @@ const OPENAI_COMPATIBLE_PROVIDERS: &[&str] = &[ "meta", "cognition", "scx-ai", + "zerogpu", ]; impl ExceptionFamily { diff --git a/litellm-rust/crates/model-catalog/src/lib.rs b/litellm-rust/crates/model-catalog/src/lib.rs index 066c2c83c6b..d42c05e6fe7 100644 --- a/litellm-rust/crates/model-catalog/src/lib.rs +++ b/litellm-rust/crates/model-catalog/src/lib.rs @@ -1,3 +1,36 @@ +//! The model catalog: per-model pricing, context limits and capabilities. +//! +//! This is the Rust form of `model_prices_and_context_window.json`, the map Python reads as +//! `litellm.model_cost` and through `litellm.get_model_info`. It answers "what does this model +//! cost and what can it do"; it knows nothing about configured deployments. The router's view of +//! deployments (groups, aliases, wildcard patterns, team mappings) is the router registry, defined +//! in `litellm-router-core`, which reads this catalog through its dynamic config. +//! +//! # Structure +//! - `Catalog`: the parsed map. `Catalog::parse` takes the JSON body and a `Provenance` (where +//! the body came from: the bundled copy in the repo or the remote copy, with revision and etag), +//! and keeps every entry in file order with its raw fields. +//! - `ModelEntry`: one model. `ModelEntry::fields` is the raw JSON object, so unknown fields are +//! never lost; `ModelEntry::info` is the typed mirror `ModelInfo` (pricing in `pricing`, +//! modalities and modes in `capabilities`). +//! - Aliases: an entry may list `aliases`; each alias resolves to its canonical key. Problems +//! found while indexing them (bad lists, collisions) are kept as `AliasIssue`s instead of +//! failing the parse. +//! - `FallbackGeneralizations` / `FallbackRule`: regex rules that map unknown model ids onto known +//! families. They are data in the file, not entries. +//! - Validation: `Catalog::validate` checks a catalog against `IntegrityLimits` (a remote copy that +//! lost most of its models is refused), and `validate_registry` checks every entry against the +//! checked-in contract. +//! +//! # Lookups +//! `Catalog::lookup` resolves a key exactly first (entry or alias), then case-insensitively, and +//! returns a `ModelMatch` with both the key that matched and the canonical key of the entry. Provider +//! prefixes are not stripped here; callers decide which keys to try. +//! +//! # Readers +//! Cost tracking, the router (through `RouterDynamicConfig`) and the gateway. Nothing in this crate +//! performs I/O; loading and refreshing the file is the caller's job. + mod capabilities; mod catalog; mod error; diff --git a/litellm-rust/crates/python-bridge/src/logger/tests.rs b/litellm-rust/crates/python-bridge/src/logger/tests.rs index 1fca3be720e..9a07428808a 100644 --- a/litellm-rust/crates/python-bridge/src/logger/tests.rs +++ b/litellm-rust/crates/python-bridge/src/logger/tests.rs @@ -187,6 +187,7 @@ fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication locals .set_item("http_warning", wrap_pyfunction!(http_warning, py).unwrap()) .unwrap(); + locals.set_item("native", crate::native_module(py)).unwrap(); let importable = py .eval( c"__import__('importlib.util', fromlist=['util']).find_spec('dotenv') is not None", @@ -205,6 +206,7 @@ import asyncio import logging import sys sys.path.insert(0, repo_root) +sys.modules['litellm.rust_bridge._native'] = native import litellm from litellm._logging import verbose_logger, session_id_var, trace_id_var diff --git a/litellm-rust/crates/router-core/Cargo.toml b/litellm-rust/crates/router-core/Cargo.toml new file mode 100644 index 00000000000..af97bd52c57 --- /dev/null +++ b/litellm-rust/crates/router-core/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "litellm-router-core" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "The router core behind litellm.Router: registry, settings, health, selection and the call loop" + +[dependencies] diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs new file mode 100644 index 00000000000..05d95d02c84 --- /dev/null +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -0,0 +1,63 @@ +//! Routing for LiteLLM: decide which deployment serves a request, and drive retries, fallbacks +//! and cooldowns around the attempts. +//! +//! This crate defines the boundary `litellm.Router` is moving behind. Today the Python router in +//! `litellm/router.py` implements all of it; this crate is where the Rust core will implement the +//! same interface. Until then it holds this overview and nothing else, so the interface is written +//! down in code before either side changes. +//! +//! # The public surface +//! Code outside the router uses `litellm.Router` through two kinds of members and nothing else: +//! - The endpoint methods (`acompletion`, `aembedding`, `aresponses`, `aanthropic_messages`, the +//! files, batches, assistants and vector store wrappers, 37 in all). Each is a thin wrapper over +//! `call(op, request)`. +//! - Eight entry points: `registry()`, `settings()`, `health()`, `select()`, `apply()`, +//! `observe()`, `register_plugin()` and `retire()`. +//! +//! Every other member that outside code uses today has a named replacement among those. The old +//! names stay available for about a month through a `DeprecatedRouterMembers` mixin, which +//! implements them over the entry points and warns on use. A CI check +//! (`tests/code_coverage_tests/router_surface_check.py`) keeps an allowlist of the members outside +//! code touches; the list may only shrink. +//! +//! # Structure +//! - `Registry`: the configured deployments, indexed by id, model group, alias, wildcard pattern +//! and team. Immutable and versioned. Seven queries: `deployment(id)`, `deployments(filter)`, +//! `names(filter)`, `resolve(name, scope)`, `group(name)`, `access_groups()` and +//! `credentials(id)`. "Registry" is the router's view of what is configured; per-model pricing +//! and capabilities are the model catalog (`litellm-model-catalog`), which the router reads +//! through its dynamic config. +//! - `RouterConfig` / `RouterDynamicConfig`: everything the router needs to know, passed in. The +//! router never reads process globals. Static config (model list, strategy, retry, timeout and +//! cooldown policy, Redis settings, plugins) is fixed for the router's life. Dynamic config +//! (pricing, model aliases, fallbacks, retry counts, callbacks) is a versioned snapshot read +//! once per call. +//! - `Settings`: one typed, frozen view of the configuration a caller may read. +//! - `Health`: cooldowns, blocked deployments and usage, kept in shared cache keys so Python and +//! Rust workers see the same state. `report(id, outcome)` records outcomes the router did not +//! see itself (a health check, an interceptor that marks a deployment failed). +//! - `select(model, request)`: resolve a name, filter eligible deployments, pick one with the +//! configured strategy, without calling it. +//! - `apply(changes)`: one `ChangeSet` (upserts, deletes, settings), applied atomically, producing +//! a new registry version. +//! - `call(op, request)`: select, run the attempt, and retry or fall back until it succeeds or the +//! policy is exhausted. +//! +//! # Execution +//! Each attempt resolves an `Executor` after a deployment is picked. A Rust-native operation runs +//! its `inference-*` route machine in-process; anything else goes through `Host::invoke`. A +//! started attempt is never replayed on a different executor. +//! +//! # Hosts +//! `Host` is what the embedding process provides: `invoke` for non-native operations, `select` +//! for strategies and sub-routers not yet ported, `allow_fallback` for access and budget checks, +//! `sleep`, `observe` for routing events, and `cache`, a DualCache-style store that holds all +//! routing state under the same keys, encodings and TTLs the Python router uses. `PythonHost` +//! serves them from the existing Python code when the router runs inside Python; `NativeHost` +//! serves them in the Rust server, where an operation without a Rust implementation makes the +//! config fail at build time, never partway through a request. +//! +//! # Plugins +//! Routing strategies, pre-call checks and sub-routers (auto, complexity, quality, adaptive) +//! register through `register_plugin` and see a `RoutingContext` (registry, settings, health and +//! the current request), never the router itself. diff --git a/litellm/__init__.py b/litellm/__init__.py index af771b2f0e6..c649a48211a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -681,6 +681,7 @@ moonshot_models: Set = set() publicai_models: Set = set() darkbloom_models: Set = set() openinfer_models: Final[Set[str]] = set() +zerogpu_models: Final[Set[str]] = set() v0_models: Set = set() morph_models: Set = set() opencode_go_models: Final[set[str]] = set() # mutable-ok: populated when the shared model catalog loads @@ -969,6 +970,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje darkbloom_models.add(key) elif value.get("litellm_provider") == "openinfer": openinfer_models.add(key) + elif value.get("litellm_provider") == "zerogpu": + zerogpu_models.add(key) elif value.get("litellm_provider") == "v0": v0_models.add(key) elif value.get("litellm_provider") == "morph": @@ -1142,6 +1145,7 @@ model_list = list( | publicai_models | darkbloom_models | openinfer_models + | zerogpu_models | v0_models | morph_models | opencode_go_models @@ -1261,6 +1265,7 @@ def _build_models_by_provider() -> dict: "publicai": publicai_models, "darkbloom": darkbloom_models, "openinfer": openinfer_models, + "zerogpu": zerogpu_models, "v0": v0_models, "morph": morph_models, "opencode_go": opencode_go_models, @@ -1331,7 +1336,6 @@ all_embedding_models = ( ) ####### IMAGE GENERATION MODELS ################### -openai_image_generation_models = ["dall-e-2", "dall-e-3"] ####### VIDEO GENERATION MODELS ################### openai_video_generation_models = ["sora-2"] @@ -1739,6 +1743,9 @@ if TYPE_CHECKING: from .llms.lm_studio.decisions.transformation import ( LMStudioDecisionsConfig as LMStudioDecisionsConfig, ) + from .llms.ollama.decisions.transformation import ( + OllamaDecisionsConfig as OllamaDecisionsConfig, + ) from .llms.laya.decisions.transformation import ( LayaDecisionsConfig as LayaDecisionsConfig, ) diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index f43e369e092..599234bab1a 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -171,6 +171,7 @@ LLM_CONFIG_NAMES: Final = ( "DatabricksDecisionsConfig", "HostedVLLMDecisionsConfig", "LMStudioDecisionsConfig", + "OllamaDecisionsConfig", "LayaDecisionsConfig", "BespokeDecisionsConfig", "OpenAIDecisionsConfig", @@ -749,6 +750,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { "HostedVLLMDecisionsConfig", ), "LMStudioDecisionsConfig": (".llms.lm_studio.decisions.transformation", "LMStudioDecisionsConfig"), + "OllamaDecisionsConfig": (".llms.ollama.decisions.transformation", "OllamaDecisionsConfig"), "LayaDecisionsConfig": (".llms.laya.decisions.transformation", "LayaDecisionsConfig"), "BespokeDecisionsConfig": (".llms.bespoke.decisions.transformation", "BespokeDecisionsConfig"), "OpenAIDecisionsConfig": (".llms.openai.decisions.transformation", "OpenAIDecisionsConfig"), diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py index 1cea0194715..d4400fbb4cc 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py @@ -19,6 +19,7 @@ from litellm.a2a_protocol.utils import ( ) from litellm.exceptions import BadRequestError from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines RUNTIME_SESSION_ID_MIN_LENGTH: Final = 33 RUNTIME_SESSION_ID_MAX_LENGTH: Final = 256 @@ -50,7 +51,7 @@ _RESERVED_PREFIX_HEADERS: Final[tuple[str, ...]] = ( class _SSELineSource(Protocol): """Minimal streaming-response surface used to read SSE lines.""" - def aiter_lines(self) -> AsyncIterator[str]: ... + def aiter_text(self) -> AsyncIterator[str]: ... def _filter_reserved_headers( @@ -244,7 +245,7 @@ class BedrockAgentCoreA2ATransformation: Yields: Parsed JSON dicts from SSE data lines """ - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): line = line.strip() if not line: continue diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index 0a2871d25cc..3937c47ea3c 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.custom_http import httpxSpecialProvider _IBM_CLOUD_IAM_URL: Final = "https://iam.cloud.ibm.com/identity/token" @@ -74,7 +75,7 @@ class _WXORun(TypedDict, total=False): class _SSELineSource(Protocol): - def aiter_lines(self) -> AsyncIterator[str]: ... + def aiter_text(self) -> AsyncIterator[str]: ... class _WXOView(TypedDict, total=False): @@ -231,7 +232,7 @@ class WatsonxOrchestrateHandler: async def _accumulate_wxo_sse_text(response: _SSELineSource) -> str: source: Final[_WXOView] = {"sse_source": response} accumulated_text = "" - async for line in source["sse_source"].aiter_lines(): + async for line in aiter_sse_lines(source["sse_source"].aiter_text()): if not line.startswith("data:"): continue data_str = line[5:].strip() diff --git a/litellm/accounting/__init__.py b/litellm/accounting/__init__.py new file mode 100644 index 00000000000..d4e56f61279 --- /dev/null +++ b/litellm/accounting/__init__.py @@ -0,0 +1,62 @@ +"""Public accounting contracts. Importing these values does not enable accounting.""" + +from litellm.accounting.measurements import measurements_for +from litellm.accounting.protocols import BudgetEnforcer, CostAggregator, CostCompiler, UsageMeter +from litellm.accounting.types import ( + AdmissionDecision, + AggregateQuery, + AggregateRow, + CostEvent, + DeliveryState, + DeliveryTarget, + Dimension, + DimensionFilter, + DimensionName, + EventKind, + Granularity, + InboundUsageEvent, + Measurement, + Metric, + MetricName, + Multiplier, + NamedCost, + OutboundUsageEvent, + Outcome, + PricedQuantity, + PricingSnapshot, + Rate, + Unpriced, + validate_aggregate_rows, +) + +__all__ = [ + "AdmissionDecision", + "AggregateQuery", + "AggregateRow", + "BudgetEnforcer", + "CostAggregator", + "CostCompiler", + "CostEvent", + "DeliveryState", + "DeliveryTarget", + "Dimension", + "DimensionFilter", + "DimensionName", + "EventKind", + "Granularity", + "InboundUsageEvent", + "Measurement", + "Metric", + "MetricName", + "Multiplier", + "NamedCost", + "OutboundUsageEvent", + "Outcome", + "PricedQuantity", + "PricingSnapshot", + "Rate", + "Unpriced", + "UsageMeter", + "measurements_for", + "validate_aggregate_rows", +] diff --git a/litellm/accounting/measurements.py b/litellm/accounting/measurements.py new file mode 100644 index 00000000000..83011a280da --- /dev/null +++ b/litellm/accounting/measurements.py @@ -0,0 +1,13 @@ +from typing import Final + +from litellm.accounting.types import CostEvent, InboundUsageEvent, Measurement, OutboundUsageEvent + + +def measurements_for(event: InboundUsageEvent | OutboundUsageEvent | CostEvent) -> tuple[Measurement, ...]: + """Return the complete contribution set for one event, ready for one apply call.""" + metrics: Final = event.charges if isinstance(event, CostEvent) else event.metrics + occurred_at: Final = event.occurred_at if isinstance(event, CostEvent) else event.started_at + return tuple( + Measurement(source_event_id=event.event_id, occurred_at=occurred_at, dimensions=event.dimensions, metric=metric) + for metric in metrics or () + ) diff --git a/litellm/accounting/protocols.py b/litellm/accounting/protocols.py new file mode 100644 index 00000000000..8955f56e7a0 --- /dev/null +++ b/litellm/accounting/protocols.py @@ -0,0 +1,85 @@ +"""Contracts only. Implementations own persistence, pricing, projections and live budgets. + +The event row's ``payload`` is a private envelope versioned by the row's ``schema_version``. +Usage envelopes retain the public event and request-scoped pricing context: quantities with units, +model/provider, requested and billed tier/region when known, and per-request/deployment overrides. +These inputs must survive request teardown. Cost-map rates are resolved at compilation, not capture. +Cost envelopes retain the public CostEvent, including its required finalized pricing snapshot. +T5/T6 own envelope serialization; this public API does not introduce a second persisted schema. +""" + +from typing import Protocol + +from litellm.accounting.types import ( + AdmissionDecision, + AggregateQuery, + AggregateRow, + CostEvent, + Dimension, + InboundUsageEvent, + Measurement, + Metric, + OutboundUsageEvent, + Unpriced, +) + + +class UsageMeter(Protocol): + async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None: + """Acknowledge commit, no-op identical IDs and reject conflicting content. + + Persist the event and its request-scoped pricing context before acknowledging. + HTTP acknowledgement is a separate policy. Unknown usage remains None. + """ + ... + + +class CostCompiler(Protocol): + async def compile(self, event: OutboundUsageEvent) -> CostEvent | Unpriced: + """Reuse the committed charge or finalize one charge and its delivery intent. + + Resolve captured request-scoped context by event identity and cost-map rates + at compilation. Persist the resolved inputs in the required snapshot. Replay + reuses the committed result and never consumes free allowance again. + Missing pricing is Unpriced, never an inferred zero. + """ + ... + + +class CostAggregator(Protocol): + async def apply(self, measurements: tuple[Measurement, ...]) -> None: + """Apply one source event's complete contribution set per target/version. + + Each target commits all contributions and its receipt atomically. A receipt + for the first metric must not suppress the rest. Usage need not wait for cost. + """ + ... + + async def query(self, query: AggregateQuery) -> tuple[AggregateRow, ...]: + """Return disjoint [start, end) rows per group using existing daily projections. + + Filters combine with AND, values within a filter with IN. DAY emits daily + buckets; None emits whole-interval totals. Reject unsupported metric/scope + combinations. Different groups may overlap, including non-additive tags. + """ + ... + + +class BudgetEnforcer(Protocol): + async def reserve( + self, reservation_id: str, dimensions: tuple[Dimension, ...], estimated: tuple[Metric, ...] + ) -> AdmissionDecision: + """Apply existing admission policies and estimates; duplicate IDs add no hold.""" + ... + + async def settle(self, reservation_id: str | None, cost_event: CostEvent) -> None: + """Settle the committed charge with atomic deduplication and hold mutation. + + Preserve scope/window attribution and retry unfinished settlement after + restart. SQL projection receipts do not provide live-counter atomicity. + """ + ... + + async def release(self, reservation_id: str) -> None: + """Idempotently release unused exposure at the request boundary, preserving charges.""" + ... diff --git a/litellm/accounting/types.py b/litellm/accounting/types.py new file mode 100644 index 00000000000..d0be29bf5b9 --- /dev/null +++ b/litellm/accounting/types.py @@ -0,0 +1,428 @@ +"""Accounting values use UTC instants, decimal strings on the wire and USD charges.""" + +from collections.abc import Mapping +from datetime import datetime, timezone +from decimal import Decimal +from enum import Enum +from typing import Annotated, Final, Literal, TypeAlias + +from pydantic import AfterValidator, AwareDatetime, BaseModel, ConfigDict, Field, model_validator +from typing_extensions import Self + +from litellm.types.utils import CallTypes + +_CONFIG: Final = ConfigDict(frozen=True, extra="forbid", validate_default=True) +Identifier: TypeAlias = Annotated[str, Field(strict=True, min_length=1, pattern=r"\S")] +Quantity: TypeAlias = Annotated[Decimal, Field(strict=True, ge=0, allow_inf_nan=False)] + + +def _utc(value: datetime) -> datetime: + return value.astimezone(timezone.utc) + + +Instant: TypeAlias = Annotated[AwareDatetime, AfterValidator(_utc)] + + +class MetricName(str, Enum): + """Reasoning tokens are included in completion tokens; above-1hr writes are included in cache creation tokens.""" + + REQUESTS = "requests" + SUCCESSFUL_REQUESTS = "successful_requests" + FAILED_REQUESTS = "failed_requests" + PROMPT_TOKENS = "prompt_tokens" + COMPLETION_TOKENS = "completion_tokens" + REASONING_TOKENS = "reasoning_tokens" + TOTAL_TOKENS = "total_tokens" + CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens" + CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens" + CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR = "cache_creation_input_tokens_above_1hr" + SPEND = "spend" + + +class DimensionName(str, Enum): + KEY = "key" + USER = "user" + TEAM = "team" + ORGANIZATION = "organization" + END_USER = "end_user" + MODEL = "model" + MODEL_GROUP = "model_group" + DEPLOYMENT = "deployment" + PROVIDER = "provider" + TOOL = "tool" + AGENT = "agent" + PROJECT = "project" + MODEL_ACCESS_GROUP = "model_access_group" + API_BASE = "api_base" + SERVICE_TIER = "service_tier" + TAG = "tag" + + +class Outcome(str, Enum): + SUCCEEDED = "succeeded" + FAILED = "failed" + CANCELLED = "cancelled" + + +class Granularity(str, Enum): + DAY = "day" + + +class EventKind(str, Enum): + INBOUND_USAGE = "inbound_usage" + OUTBOUND_USAGE = "outbound_usage" + COST = "cost" + + +class DeliveryState(str, Enum): + PENDING = "pending" + APPLIED = "applied" + PARKED = "parked" + + +class DeliveryTarget(str, Enum): + """One member per existing table the projector writes; each is applied and receipted independently.""" + + KEY_SPEND = "LiteLLM_VerificationToken" + USER_SPEND = "LiteLLM_UserTable" + TEAM_SPEND = "LiteLLM_TeamTable" + TEAM_MEMBER_SPEND = "LiteLLM_TeamMembership" + ORGANIZATION_SPEND = "LiteLLM_OrganizationTable" + ORGANIZATION_MEMBER_SPEND = "LiteLLM_OrganizationMembership" + END_USER_SPEND = "LiteLLM_EndUserTable" + AGENT_SPEND = "LiteLLM_AgentsTable" + TAG_SPEND = "LiteLLM_TagTable" + PROJECT_SPEND = "LiteLLM_ProjectTable" + MODEL_ACCESS_GROUP_SPEND = "LiteLLM_ModelAccessGroupBudgetTable" + BUDGET_SPEND = "LiteLLM_BudgetTable" + BUDGET_WINDOW_SPEND = "LiteLLM_BudgetWindowSpend" + DAILY_USER_SPEND = "LiteLLM_DailyUserSpend" + DAILY_TEAM_SPEND = "LiteLLM_DailyTeamSpend" + DAILY_ORGANIZATION_SPEND = "LiteLLM_DailyOrganizationSpend" + DAILY_END_USER_SPEND = "LiteLLM_DailyEndUserSpend" + DAILY_AGENT_SPEND = "LiteLLM_DailyAgentSpend" + DAILY_TAG_SPEND = "LiteLLM_DailyTagSpend" + DAILY_MODEL_USAGE = "LiteLLM_DailyModelUsage" + DAILY_TOOL_SPEND = "LiteLLM_DailyToolSpend" + AUTOROUTER_DAILY_SPEND = "LiteLLM_AutoRouterDailySpend" + + +_SUBSET_PAIRS: Final = ( + (MetricName.REASONING_TOKENS, MetricName.COMPLETION_TOKENS), + (MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, MetricName.CACHE_CREATION_INPUT_TOKENS), +) + + +def _check_subsets(values: Mapping[str, Decimal]) -> None: + for part, whole in _SUBSET_PAIRS: + if part.value in values and whole.value in values and values[part.value] > values[whole.value]: + raise ValueError(f"{part.value} cannot exceed {whole.value}") + + +class Metric(BaseModel): + model_config = _CONFIG + + name: MetricName + value: Quantity + unit: Literal["count", "USD"] = "count" + + @model_validator(mode="after") + def validate_unit(self) -> Self: + if (self.name is MetricName.SPEND) != (self.unit == "USD"): + raise ValueError("SPEND requires USD; usage metrics require count") + if self.unit == "count" and self.value != self.value.to_integral_value(): + raise ValueError("Count metrics require integral values") + return self + + +class Dimension(BaseModel): + model_config = _CONFIG + + name: DimensionName + value: Identifier + + +_REPEATABLE_DIMENSIONS: Final = frozenset({DimensionName.TAG, DimensionName.MODEL_ACCESS_GROUP}) + + +def _dimensions(values: tuple[Dimension, ...]) -> tuple[Dimension, ...]: + if len(frozenset(values)) != len(values): + raise ValueError("Duplicate dimension values") + single_names: Final = tuple(value.name for value in values if value.name not in _REPEATABLE_DIMENSIONS) + if len(frozenset(single_names)) != len(single_names): + raise ValueError("Only tag and model access group dimensions may have multiple values") + return values + + +Dimensions: TypeAlias = Annotated[tuple[Dimension, ...], AfterValidator(_dimensions)] +_CALLER_METRICS: Final = frozenset((MetricName.REQUESTS, MetricName.SUCCESSFUL_REQUESTS, MetricName.FAILED_REQUESTS)) + + +def _metrics(values: tuple[Metric, ...]) -> tuple[Metric, ...]: + if len(frozenset(value.name for value in values)) != len(values): + raise ValueError("Duplicate metric names") + return values + + +Metrics: TypeAlias = Annotated[tuple[Metric, ...], AfterValidator(_metrics)] + + +def _interval(start: datetime, end: datetime, *, allow_empty: bool = False) -> None: + if end < start or (end == start and not allow_empty): + raise ValueError("End must follow start") + + +class InboundUsageEvent(BaseModel): + model_config = _CONFIG + + event_id: Identifier + request_id: Identifier + call_type: CallTypes + started_at: Instant + ended_at: Instant + outcome: Outcome + dimensions: Dimensions + metrics: Metrics + + @model_validator(mode="after") + def validate_caller_activity(self) -> Self: + _interval(self.started_at, self.ended_at, allow_empty=True) + if any(metric.name not in _CALLER_METRICS for metric in self.metrics): + raise ValueError("Inbound events own only caller request metrics") + if not any(metric.name is MetricName.REQUESTS and metric.value == 1 for metric in self.metrics): + raise ValueError("Inbound events count exactly one caller request") + for metric in self.metrics: + if metric.name is MetricName.SUCCESSFUL_REQUESTS and metric.value != int(self.outcome is Outcome.SUCCEEDED): + raise ValueError("Successful request count must match outcome") + if metric.name is MetricName.FAILED_REQUESTS and metric.value != int(self.outcome is not Outcome.SUCCEEDED): + raise ValueError("Failed request count must match outcome") + return self + + +class OutboundUsageEvent(BaseModel): + model_config = _CONFIG + + event_id: Identifier + attempt_id: Identifier + request_id: Identifier | None + retry_of_attempt_id: Identifier | None + call_type: CallTypes + started_at: Instant + ended_at: Instant + outcome: Outcome + dimensions: Dimensions + metrics: Metrics | None + + @model_validator(mode="after") + def validate_attempt_usage(self) -> Self: + _interval(self.started_at, self.ended_at, allow_empty=True) + if self.retry_of_attempt_id == self.attempt_id: + raise ValueError("An attempt cannot retry itself") + if self.metrics is not None and any( + metric.name in _CALLER_METRICS or metric.name is MetricName.SPEND for metric in self.metrics + ): + raise ValueError("Outbound events own only attempt usage metrics") + if self.metrics is not None: + _check_subsets({metric.name.value: metric.value for metric in self.metrics}) + return self + + +PricingUnit: TypeAlias = Literal["token", "second", "character", "image", "request"] + + +class Rate(BaseModel): + model_config = _CONFIG + + measure: Identifier + unit: PricingUnit + usd_per_unit: Quantity + + +class PricedQuantity(BaseModel): + """A pricing input in its billing unit. Only seconds may be fractional.""" + + model_config = _CONFIG + + measure: Identifier + unit: PricingUnit + value: Quantity + + @model_validator(mode="after") + def validate_discrete_quantity(self) -> Self: + if self.unit != "second" and self.value != self.value.to_integral_value(): + raise ValueError("Discrete priced quantities require integral values") + return self + + +class Multiplier(BaseModel): + """A regional or tier factor applied on top of the priced key, e.g. data_residency=eu or vertex_location=us-east5.""" + + model_config = _CONFIG + + name: Identifier + value: Identifier + factor: Annotated[Decimal, Field(strict=True, gt=0, allow_inf_nan=False)] + + +class NamedCost(BaseModel): + """A free-form extra charge by name, e.g. azure_model_router_flat_cost.""" + + model_config = _CONFIG + + name: Identifier + amount: Quantity + + +class PricingSnapshot(BaseModel): + """Retained inputs and result of authoritative compilation, reproducible after cost-map rates change. + + Request-scoped inputs and overrides come from the usage row's versioned payload. Cost-map rates are + resolved at compilation, not frozen at capture. Retries reuse the committed charge and snapshot. + Component names match ``CostBreakdown``: reasoning_cost is part of output_cost, not an additional charge. + Component sums and float tolerances belong to reconciliation, not this structural contract. + """ + + model_config = _CONFIG + + schema_version: Literal[1] = 1 + calculator_version: Identifier + model: Identifier + provider: Identifier | None + pricing_source: Literal["cost_map", "deployment", "per_request", "upstream_header"] + priced_model_key: Identifier + multipliers: tuple[Multiplier, ...] = () + rates: tuple[Rate, ...] = () + usage: tuple[PricedQuantity, ...] = () + input_cost: Quantity | None = None + output_cost: Quantity | None = None + cache_read_cost: Quantity | None = None + cache_creation_cost: Quantity | None = None + reasoning_cost: Quantity | None = None + tool_usage_cost: Quantity | None = None + guardrail_cost: Quantity | None = None + additional_costs: tuple[NamedCost, ...] = () + original_cost: Quantity | None = None + discount_percent: Quantity | None = None + discount_amount: Quantity | None = None + margin_percent: Quantity | None = None + margin_fixed_amount: Quantity | None = None + margin_total_amount: Quantity | None = None + upstream_reported_cost: Quantity | None = None + total_cost: Quantity + + @model_validator(mode="after") + def validate_snapshot(self) -> Self: + rate_keys: Final = frozenset((rate.measure, rate.unit) for rate in self.rates) + quantity_keys: Final = frozenset((quantity.measure, quantity.unit) for quantity in self.usage) + if len(rate_keys) != len(self.rates): + raise ValueError("Duplicate rate measures and units") + if len(quantity_keys) != len(self.usage): + raise ValueError("Duplicate priced quantity measures and units") + if not rate_keys.issubset(quantity_keys): + raise ValueError("Every rate requires one priced quantity with the same measure and unit") + if len(frozenset(cost.name for cost in self.additional_costs)) != len(self.additional_costs): + raise ValueError("Duplicate additional cost names") + if len(frozenset(multiplier.name for multiplier in self.multipliers)) != len(self.multipliers): + raise ValueError("Duplicate multiplier names") + _check_subsets({quantity.measure: quantity.value for quantity in self.usage if quantity.unit == "token"}) + return self + + +class CostEvent(BaseModel): + model_config = _CONFIG + + event_id: Identifier + usage_event_id: Identifier + occurred_at: Instant + dimensions: Dimensions + charges: Metrics + snapshot: PricingSnapshot + + @model_validator(mode="after") + def validate_charge(self) -> Self: + if len(self.charges) != 1 or self.charges[0].name is not MetricName.SPEND: + raise ValueError("A cost event requires one authoritative USD spend metric") + if self.event_id == self.usage_event_id: + raise ValueError("Cost and usage events require distinct identities") + if self.snapshot.total_cost != self.charges[0].value: + raise ValueError("Snapshot total must equal the authoritative spend") + return self + + +class Unpriced(BaseModel): + model_config = _CONFIG + + usage_event_id: Identifier + reason: Identifier + + +class Measurement(BaseModel): + model_config = _CONFIG + + source_event_id: Identifier + occurred_at: Instant + dimensions: Dimensions + metric: Metric + + +class DimensionFilter(BaseModel): + model_config = _CONFIG + + name: DimensionName + values: Annotated[tuple[Identifier, ...], Field(min_length=1)] + + +class AggregateQuery(BaseModel): + model_config = _CONFIG + + metric: MetricName + filters: tuple[DimensionFilter, ...] + group_by: tuple[DimensionName, ...] + start: Instant + end: Instant + granularity: Granularity | None + + @model_validator(mode="after") + def validate_query(self) -> Self: + _interval(self.start, self.end) + if len(frozenset(item.name for item in self.filters)) != len(self.filters): + raise ValueError("Use one IN filter per dimension; filters combine with AND") + if len(frozenset(self.group_by)) != len(self.group_by): + raise ValueError("Duplicate group-by dimensions") + if any( + (instant.hour, instant.minute, instant.second, instant.microsecond) != (0, 0, 0, 0) + for instant in (self.start, self.end) + ): + raise ValueError("Initial reporting queries require UTC day boundaries") + return self + + +class AggregateRow(BaseModel): + model_config = _CONFIG + + start: Instant + end: Instant + dimensions: Dimensions + value: Quantity + + @model_validator(mode="after") + def validate_bounds(self) -> Self: + _interval(self.start, self.end) + return self + + +def validate_aggregate_rows(rows: tuple[AggregateRow, ...]) -> tuple[AggregateRow, ...]: + """Validate [start, end) buckets per group; overlapping tag groups are not additive.""" + groups: Final = frozenset(frozenset(row.dimensions) for row in rows) + for group in groups: + ordered = sorted((row for row in rows if frozenset(row.dimensions) == group), key=lambda row: row.start) + if any(left.end > right.start for left, right in zip(ordered, ordered[1:])): + raise ValueError("Aggregate buckets overlap within a dimension group") + return rows + + +class AdmissionDecision(BaseModel): + model_config = _CONFIG + + allowed: bool + reason: Identifier | None = None diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 1cc3e106aa0..6b40c48c4ef 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -274,7 +274,9 @@ class InMemoryCache(BaseCache): ) -> list[float] | None: results: Final = [] for increment in increment_list: - result = await self.async_increment(increment["key"], increment["increment_value"], **kwargs) + result = await self.async_increment( + increment["key"], increment["increment_value"], ttl=increment["ttl"], **kwargs + ) results.append(result) return results diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 98322aecf6a..ee3639ac387 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -3,7 +3,6 @@ Handler for transforming /chat/completions api requests to litellm.responses req """ import json -import os from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from itertools import accumulate, chain from types import MappingProxyType @@ -35,6 +34,7 @@ from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.bridges.completion_transformation import ( CompletionTransformationBridge, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.responses.litellm_completion_transformation.tool_search import is_anthropic_tool_search_chat_tool from litellm.responses.sse_output_recovery import ( parse_sse_json_chunk, @@ -43,8 +43,8 @@ from litellm.responses.sse_output_recovery import ( ) from litellm.responses.utils import ( ResponsesAPIRequestUtils, - auto_reasoning_summary, normalize_responses_api_stream_options, + opt_in_auto_reasoning_summary, ) from litellm.types.llms.openai import ( ChatCompletionAnnotation, @@ -1352,7 +1352,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): recovered_output_items: Final[dict[int, dict[str, object]]] = {} recovered_text_only_items: Final[dict[int, dict[str, object]]] = {} - for chunk in raw_sse.splitlines(): + for chunk in split_sse_lines(raw_sse): parsed_chunk = parse_sse_json_chunk(chunk) if parsed_chunk is None: continue @@ -1745,19 +1745,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): **cast(Reasoning, reasoning_effort) # cast-ok: dict is forwarded verbatim to the provider ) - # Check if auto-summary is enabled via flag or environment variable - # Priority: litellm.reasoning_auto_summary flag > LITELLM_REASONING_AUTO_SUMMARY env var - auto_summary_enabled: Final = ( - litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true" - ) - + summary: Final = opt_in_auto_reasoning_summary(model, custom_llm_provider) return ( - Reasoning( - effort=reasoning_effort, - summary="detailed" if model is None else auto_reasoning_summary(model, custom_llm_provider), - ) - if auto_summary_enabled - else Reasoning(effort=reasoning_effort) + Reasoning(effort=reasoning_effort) + if summary is None + else Reasoning(effort=reasoning_effort, summary=summary) ) _map_reasoning_effort = map_reasoning_effort diff --git a/litellm/constants.py b/litellm/constants.py index 496f82b3264..81b3e7c2f8d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -130,6 +130,7 @@ MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024 S3_LOG_PROMPTS_ONLY_ENV_VAR: Final = "S3_LOG_PROMPTS_ONLY" S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY" MAX_FILE_LIST_LIMIT: Final = 10000 +ANTHROPIC_DEFAULT_FILE_LIST_LIMIT: Final = 20 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10) DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = get_env_int("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1) GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10 @@ -764,7 +765,6 @@ ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = { # Used for web search interception across providers LITELLM_WEB_SEARCH_TOOL_NAME: Final = "litellm_web_search" -DEFAULT_IMAGE_ENDPOINT_MODEL: Final = "dall-e-2" DEFAULT_VIDEO_ENDPOINT_MODEL: Final = "sora-2" DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS: Final = get_env_int("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8) @@ -1097,6 +1097,7 @@ openai_compatible_providers: Final[list[str]] = [ "llamafile", "lm_studio", "galadriel", + "zerogpu", "github_copilot", # GitHub Copilot Chat API "chatgpt", # ChatGPT subscription API "novita", diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 86b22069a39..087d5b5b718 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1592,14 +1592,6 @@ def completion_cost( for part in ws_tier_parts ) - if ( - (call_type == "aimage_generation" or call_type == "image_generation") - and model is not None - and isinstance(model, str) - and len(model) == 0 - and custom_llm_provider == "azure" - ): - model = "dall-e-2" # for dall-e-2, azure expects an empty model name # Handle Inputs to completion_cost prompt_tokens = 0 prompt_characters: int | None = None @@ -1674,7 +1666,8 @@ def completion_cost( raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates) ) - for idx, (model, candidate_pricing_provider) in enumerate(potential_model_names): + for idx, (candidate_model, candidate_pricing_provider) in enumerate(potential_model_names): + model = candidate_model try: if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug("selected model name for cost calculation: %s", model) diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index a22c47cb1e7..08ac7d4f9dc 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) @@ -131,9 +132,9 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent self.model = model self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata - # Gemini streamGenerateContent uses SSE line framing; iter_lines keeps + # Gemini streamGenerateContent uses SSE line framing; iter_sse_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. - self.stream_iterator = response.iter_lines() + self.stream_iterator = iter_sse_lines(response.iter_text()) def __iter__(self): return self @@ -182,9 +183,9 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo self.model = model self.generate_content_provider_config = generate_content_provider_config self.litellm_metadata = litellm_metadata - # Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps + # Gemini streamGenerateContent uses SSE line framing; aiter_sse_lines keeps # large inlineData payloads (e.g. image/jpeg) intact within one event. - self.stream_iterator = response.aiter_lines() + self.stream_iterator = aiter_sse_lines(response.aiter_text()) def __aiter__(self): return self diff --git a/litellm/images/main.py b/litellm/images/main.py index 84156919774..6e77cbb7046 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -1,9 +1,9 @@ import asyncio import contextvars import importlib -from collections.abc import Coroutine, Mapping +from collections.abc import Collection, Coroutine, Mapping from functools import partial -from typing import TYPE_CHECKING, Final, Literal, cast, overload +from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload from pydantic import TypeAdapter, ValidationError @@ -16,7 +16,6 @@ import litellm # client is imported from litellm as it's a decorator from litellm import client -from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT from litellm.exceptions import LiteLLMUnknownProvider from litellm.litellm_core_utils.get_llm_provider_logic import ( @@ -61,6 +60,7 @@ from litellm.types.utils import ( ProviderRoutingHints, ) from litellm.utils import ( + EXTRA_BODY_ROUTING_KEYS, ImageResponse, ProviderConfigManager, get_llm_provider, @@ -298,17 +298,14 @@ def image_generation( if custom_llm_provider == "chatgpt" or (model is not None and model.startswith("chatgpt/")) else None ) - dynamic_api_key: str | None = None - if model is not None or custom_llm_provider is not None: - model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( - model=model, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - litellm_params=chatgpt_params, - ) - else: - model = "dall-e-2" - custom_llm_provider = "openai" # default to dall-e-2 on openai + if model is None and custom_llm_provider is None: + raise litellm.BadRequestError(message="model is required", model="", llm_provider="") + model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( + model=model, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + litellm_params=chatgpt_params, + ) model_response.hidden_params["model"] = model openai_params: Final = [ "user", @@ -466,7 +463,7 @@ def image_generation( litellm.LlmProviders.QWENCLOUD, litellm.LlmProviders.QWEN_AI_PLATFORM, litellm.LlmProviders.EDENAI, - ): + ) or (custom_llm_provider == litellm.LlmProviders.TENCENT and image_generation_config is not None): if image_generation_config is None: raise ValueError(f"image generation config is not supported for {custom_llm_provider}") @@ -711,7 +708,7 @@ async def aimage_variation(*args, **kwargs) -> ImageResponse: @client def image_variation( image: FileTypes, - model: str = "dall-e-2", # set to dall-e-2 by default - like OpenAI. + model: str | None = None, n: int = 1, response_format: Literal["url", "b64_json"] = "url", size: str | None = None, @@ -727,7 +724,7 @@ def image_variation( litellm_params: Final = get_litellm_params(**kwargs) # get the custom llm provider model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider( - model=model, + model=model or "", custom_llm_provider=litellm_params.get("custom_llm_provider", None), api_base=litellm_params.get("api_base", None), api_key=litellm_params.get("api_key", None), @@ -746,7 +743,7 @@ def image_variation( response: ImageResponse | None = None provider_config: Final = ProviderConfigManager.get_provider_model_info( - model=model or "", # openai defaults to dall-e-2 + model=model, provider=llm_provider, ) @@ -771,7 +768,7 @@ def image_variation( model_response=model_response, api_key=api_key, api_base=api_base, - model=model, + model=model or None, image=image, timeout=litellm_params.get("timeout", None), custom_llm_provider=custom_llm_provider, @@ -789,7 +786,7 @@ def image_variation( model_response=model_response, api_key=api_key, api_base=api_base, - model=model, + model=model or None, image=image, timeout=litellm_params.get("timeout", None) or DEFAULT_REQUEST_TIMEOUT, custom_llm_provider=custom_llm_provider, @@ -807,6 +804,32 @@ def image_variation( return response +_IMAGE_EDIT_OPENAI_PARAMS: Final = frozenset(get_type_hints(ImageEditOptionalRequestParams)) +_DROP_PARAM_KWARGS: Final = frozenset({"drop_params", "additional_drop_params"}) +_ADDITIONAL_DROP_PARAMS: Final = TypeAdapter[tuple[str, ...]](tuple[str, ...]) + + +def _listed_drop_params(additional_drop_params: object) -> tuple[str, ...]: + try: + return _ADDITIONAL_DROP_PARAMS.validate_python(additional_drop_params) + except ValidationError: + return () + + +def _tokenhub_passthrough_edit_params( + non_default_params: Mapping[str, object], + extra_body: Mapping[str, object] | None, + mapped_params: Collection[str], + additional_drop_params: Collection[str] | None, +) -> Mapping[str, object]: + dropped: Final = frozenset(additional_drop_params or ()) + unmapped_excluded: Final = _DROP_PARAM_KWARGS | _IMAGE_EDIT_OPENAI_PARAMS | frozenset(mapped_params) | dropped + return { + **{key: value for key, value in non_default_params.items() if key not in unmapped_excluded}, + **{key: value for key, value in (extra_body or {}).items() if key not in dropped | EXTRA_BODY_ROUTING_KEYS}, + } + + @client def image_edit( image: FileTypes | list[FileTypes] | None = None, @@ -874,7 +897,7 @@ def image_edit( # get llm provider logic litellm_params: Final = GenericLiteLLMParams(**kwargs) model, custom_llm_provider, _, _ = get_llm_provider( - model=model or DEFAULT_IMAGE_ENDPOINT_MODEL, + model=model or "", custom_llm_provider=custom_llm_provider, ) @@ -935,22 +958,24 @@ def image_edit( raise ValueError(f"image edit is not supported for {custom_llm_provider}") local_vars.update(kwargs) + provider_supported_params: Final[frozenset[str]] = frozenset( + image_edit_provider_config.get_supported_openai_params(model) + ) # Get ImageEditOptionalRequestParams with only valid parameters image_edit_optional_params: Final[ImageEditOptionalRequestParams] = ( _get_ImageEditRequestUtils().get_requested_image_edit_optional_param( local_vars, - provider_supported_params=frozenset( - image_edit_provider_config.get_supported_openai_params(model) - ).intersection(non_default_params), + provider_supported_params=provider_supported_params.intersection(non_default_params), ) ) # Get optional parameters for the responses API + additional_drop_params: Final = _listed_drop_params(kwargs.get("additional_drop_params")) image_edit_request_params: Final[dict] = _get_ImageEditRequestUtils().get_optional_params_image_edit( model=model, image_edit_provider_config=image_edit_provider_config, image_edit_optional_params=image_edit_optional_params, drop_params=kwargs.get("drop_params"), - additional_drop_params=kwargs.get("additional_drop_params"), + additional_drop_params=additional_drop_params, ) if image_edit_provider_config.use_multipart_form_data() and ( @@ -1039,7 +1064,19 @@ def image_edit( image=images, prompt=prompt, image_edit_provider_config=image_edit_provider_config, - image_edit_optional_request_params=image_edit_request_params, + image_edit_optional_request_params={ + **image_edit_request_params, + **( + _tokenhub_passthrough_edit_params( + non_default_params, + extra_body if isinstance(extra_body, dict) else None, + provider_supported_params, + additional_drop_params, + ) + if custom_llm_provider == litellm.LlmProviders.TENCENT.value + else {} + ), + }, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, logging_obj=litellm_logging_obj, diff --git a/litellm/images/utils.py b/litellm/images/utils.py index 24454954714..4e0cde24a68 100644 --- a/litellm/images/utils.py +++ b/litellm/images/utils.py @@ -1,6 +1,9 @@ +import base64 +import os from collections.abc import Collection, Mapping from io import BufferedReader, BytesIO -from typing import Any, Final, cast, get_type_hints +from pathlib import Path +from typing import Any, Final, Protocol, cast, get_type_hints, runtime_checkable import litellm from litellm.litellm_core_utils.token_counter import get_image_type @@ -16,7 +19,7 @@ class ImageEditRequestUtils: image_edit_provider_config: BaseImageEditConfig, image_edit_optional_params: ImageEditOptionalRequestParams, drop_params: bool | None = None, - additional_drop_params: list[str] | None = None, + additional_drop_params: Collection[str] | None = None, ) -> dict: """ Get optional parameters for the image edit API. @@ -145,3 +148,36 @@ class ImageEditRequestUtils: except Exception: # If anything goes wrong, default to PNG return FILE_MIME_TYPES[FileType.PNG] + + +@runtime_checkable +class _SeekableBinaryReader(Protocol): + def tell(self) -> int: ... + + def seek(self, offset: int) -> int: ... + + def read(self) -> bytes: ... + + +def _read_image_bytes(image: object) -> bytes: + if isinstance(image, bytes): + return image + if isinstance(image, tuple): + return _read_image_bytes(image[1]) + if isinstance(image, os.PathLike): + return Path(image).read_bytes() + if isinstance(image, _SeekableBinaryReader): + position: Final = image.tell() + image.seek(0) + data: Final = image.read() + image.seek(position) + return data + raise ValueError(f"Unsupported image type for image edit: {type(image).__name__}") + + +def image_to_data_url(image: object) -> str: + if isinstance(image, str): + return image + image_bytes: Final = _read_image_bytes(image) + mime_type: Final = ImageEditRequestUtils.get_image_content_type(image_bytes) + return f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}" diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index d8be607150c..4907bd4368d 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -338,6 +338,8 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_ToolTable", "LiteLLM_TelemetryReport", "LiteLLM_AccessGroupTable", + "LiteLLM_AccountingDelivery", + "LiteLLM_AccountingEvent", "LiteLLM_ClaudeCodePluginTable", "LiteLLM_MemoryTable", "LiteLLM_AdaptiveRouterState", @@ -434,6 +436,7 @@ _POSTGRES_OPERATION_BY_CALL_TYPE: Final[Mapping[str, PostgresOperation]] = Mappi "reset_budget_windows": PostgresOperation("select", None), "write_budget_windows": PostgresOperation("update", None), "roll_window_spend_row": PostgresOperation("update", None), + "open_window_spend_row": PostgresOperation("upsert", None), "seed_window_spend": PostgresOperation("select", None), "select_window_spend_rows": PostgresOperation("select", None), "commit_window_spend_updates": PostgresOperation("upsert", None), diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index a38b8bed595..6e1ba3eea08 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -97,6 +97,8 @@ WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" _RESPONSE_CONTENT_FIELD: Final = "content" +_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX: Final = "Perform a web search for the query:" + _ResponseT: Final = TypeVar("_ResponseT") @@ -363,7 +365,8 @@ class WebSearchInterceptionLogger(CustomLogger): get_last_user_message, ) - query: Final = get_last_user_message(cast(list[AllMessageValues], messages)) + last_user_message: Final = get_last_user_message(cast(list[AllMessageValues], messages)) or "" + query: Final = last_user_message.removeprefix(_CLAUDE_CODE_WEB_SEARCH_PROMPT_PREFIX).strip() if not query: return None diff --git a/litellm/interactions/__init__.py b/litellm/interactions/__init__.py index 6129cd87153..d5dd9b2cd00 100644 --- a/litellm/interactions/__init__.py +++ b/litellm/interactions/__init__.py @@ -32,11 +32,11 @@ Usage: # Create a managed agent on the provider side result = litellm.interactions.agents.create( - name="waverunner", + id="waverunner", custom_llm_provider="gemini", api_key="...", - base_agent="gemini-2.5-flash", - instructions="You are a helpful assistant.", + base_agent="antigravity-preview-09-2026", + system_instruction="You are a helpful assistant.", ) Methods: diff --git a/litellm/interactions/agents/__init__.py b/litellm/interactions/agents/__init__.py index 21fd2dc0742..a81b35c2fc7 100644 --- a/litellm/interactions/agents/__init__.py +++ b/litellm/interactions/agents/__init__.py @@ -3,11 +3,11 @@ litellm.interactions.agents Full CRUD SDK for provider-side managed agents (e.g. Gemini v1beta/agents). - litellm.interactions.agents.create(name=..., ...) + litellm.interactions.agents.create(id=..., ...) litellm.interactions.agents.list(api_key=...) - litellm.interactions.agents.get(name=..., ...) - litellm.interactions.agents.delete(name=..., ...) - litellm.interactions.agents.list_versions(name=..., ...) + litellm.interactions.agents.get(id=..., ...) + litellm.interactions.agents.delete(id=..., ...) + litellm.interactions.agents.list_versions(id=..., ...) Async counterparts: acreate, alist, aget, adelete, alist_versions """ diff --git a/litellm/interactions/agents/http_handler.py b/litellm/interactions/agents/http_handler.py index 2afedc34d36..fa1e52e6842 100644 --- a/litellm/interactions/agents/http_handler.py +++ b/litellm/interactions/agents/http_handler.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, @@ -35,7 +36,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): def create_agent( self, agents_api_config: BaseAgentsAPIConfig, - name: str, + create_request: AgentCreateRequest, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, extra_headers: dict[str, str] | None = None, @@ -47,7 +48,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): if _is_async: return self.async_create_agent( agents_api_config=agents_api_config, - name=name, + create_request=create_request, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -63,12 +64,13 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): api_base=litellm_params.get("api_base"), litellm_params=dict(litellm_params), ) - data: Final = agents_api_config.transform_create_request(name=name, litellm_params=dict(litellm_params)) - if extra_body: - data.update(extra_body) + data: Final = { + **agents_api_config.transform_create_request(create_request), + **(extra_body or {}), + } logging_obj.pre_call( - input=name, + input=create_request.id, api_key="", additional_args={ "complete_input_dict": data, @@ -85,12 +87,12 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): original_response=response.text, additional_args={"complete_input_dict": data}, ) - return agents_api_config.transform_create_response(raw_response=response, name=name) + return agents_api_config.transform_create_response(raw_response=response, name=create_request.id) async def async_create_agent( self, agents_api_config: BaseAgentsAPIConfig, - name: str, + create_request: AgentCreateRequest, litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, extra_headers: dict[str, str] | None = None, @@ -106,12 +108,13 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): api_base=litellm_params.get("api_base"), litellm_params=dict(litellm_params), ) - data: Final = agents_api_config.transform_create_request(name=name, litellm_params=dict(litellm_params)) - if extra_body: - data.update(extra_body) + data: Final = { + **agents_api_config.transform_create_request(create_request), + **(extra_body or {}), + } logging_obj.pre_call( - input=name, + input=create_request.id, api_key="", additional_args={ "complete_input_dict": data, @@ -130,7 +133,7 @@ class AgentsHTTPHandler(InteractionsHTTPHandler): original_response=response.text, additional_args={"complete_input_dict": data}, ) - return agents_api_config.transform_create_response(raw_response=response, name=name) + return agents_api_config.transform_create_response(raw_response=response, name=create_request.id) # ------------------------------------------------------------------ # # LIST # diff --git a/litellm/interactions/agents/main.py b/litellm/interactions/agents/main.py index 02b0ef4807f..5e2bfa9cea2 100644 --- a/litellm/interactions/agents/main.py +++ b/litellm/interactions/agents/main.py @@ -6,41 +6,47 @@ Usage: # Create response = litellm.interactions.agents.create( - name="waverunner", + id="waverunner", custom_llm_provider="gemini", api_key="...", - base_agent="gemini-2.5-flash", - instructions="You are a helpful assistant.", + base_agent="antigravity-preview-09-2026", + system_instruction="You are a helpful assistant.", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, ) # List response = litellm.interactions.agents.list(api_key="...", custom_llm_provider="gemini") # Get - response = litellm.interactions.agents.get(name="waverunner", api_key="...") + response = litellm.interactions.agents.get(id="waverunner", api_key="...") # Delete - result = litellm.interactions.agents.delete(name="waverunner", api_key="...") + result = litellm.interactions.agents.delete(id="waverunner", api_key="...") # List versions - result = litellm.interactions.agents.list_versions(name="waverunner", api_key="...") + result = litellm.interactions.agents.list_versions(id="waverunner", api_key="...") # Async versions: acreate, alist, aget, adelete, alist_versions """ import asyncio import contextvars -from collections.abc import Coroutine, Mapping +from collections.abc import Coroutine, Mapping, Sequence +from dataclasses import dataclass from functools import partial +from types import MappingProxyType from typing import Any, Final import httpx +from pydantic import ValidationError import litellm from litellm.interactions.agents.http_handler import agents_http_handler from litellm.interactions.agents.utils import get_provider_agents_api_config +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, @@ -89,14 +95,120 @@ def _make_logging_obj( return litellm_logging_obj +@dataclass(frozen=True, slots=True) +class _InvalidAgentRequest: + message: str + + +def _alias_conflict(field: str, value: str | None, alias: str, alias_value: str | None) -> str | None: + if value is None or alias_value is None or value == alias_value: + return None + return f"'{alias}' is an alias of '{field}', but the request gives them different values. Pass only '{field}'." + + +def _agent_id_or_error(id: str | None, name: str | None) -> str | _InvalidAgentRequest: + conflict: Final = _alias_conflict("id", id, "name", name) + if conflict is not None: + return _InvalidAgentRequest(conflict) + agent_id: Final = id if id is not None else name + if agent_id is None: + return _InvalidAgentRequest("Missing required parameter: 'id' (or its alias 'name').") + return agent_id + + +def _required_agent_id(id: str | None, name: str | None, llm_provider: str) -> str: + agent_id: Final = _agent_id_or_error(id, name) + if isinstance(agent_id, _InvalidAgentRequest): + raise litellm.BadRequestError(message=agent_id.message, model=id or name or "", llm_provider=llm_provider) + return agent_id + + # ================================================================== # # CREATE # # ================================================================== # +_CREATE_FIELD_SHAPES: Final = MappingProxyType( + { + "id": "a string", + "base_agent": "a string", + "agent_config": "an object", + "system_instruction": "a string", + "base_environment": "an environment id string or an object", + "description": "a string", + "tools": "a list of objects", + } +) + + +def _invalid_create_fields_message(invalid_fields: frozenset[str], caller_names: Mapping[str, str]) -> str: + problems: Final = "; ".join( + f"'{caller_names.get(field, field)}' must be {_CREATE_FIELD_SHAPES[field]}" for field in sorted(invalid_fields) + ) + drop_hint: Final = "" if "id" in invalid_fields else " Set drop_params to drop invalid optional fields." + return f"Invalid agent create request: {problems}.{drop_hint}" + + +def _validated_create_request( + fields: Mapping[str, object], drop_invalid_fields: bool, caller_names: Mapping[str, str] +) -> AgentCreateRequest | _InvalidAgentRequest: + try: + return AgentCreateRequest.model_validate(fields) + except ValidationError as e: + invalid_fields: Final = frozenset(str(error["loc"][0]) for error in e.errors()) + if drop_invalid_fields and "id" not in invalid_fields: + return AgentCreateRequest.model_validate( + {field: value for field, value in fields.items() if field not in invalid_fields} + ) + return _InvalidAgentRequest(_invalid_create_fields_message(invalid_fields, caller_names)) + + +def _create_request_or_error( + id: str | None, + name: str | None, + base_agent: str | None, + agent_config: Mapping[str, object] | None, + system_instruction: str | None, + instructions: str | None, + base_environment: InteractionEnvironment | None, + description: str | None, + tools: Sequence[Mapping[str, object]] | None, + drop_invalid_fields: bool, +) -> AgentCreateRequest | _InvalidAgentRequest: + agent_id: Final = _agent_id_or_error(id, name) + if isinstance(agent_id, _InvalidAgentRequest): + return agent_id + instruction_conflict: Final = _alias_conflict( + "system_instruction", system_instruction, "instructions", instructions + ) + if instruction_conflict is not None and not drop_invalid_fields: + return _InvalidAgentRequest(instruction_conflict) + sent_only_instructions: Final = system_instruction is None and instructions is not None + return _validated_create_request( + MappingProxyType( + { + "id": agent_id, + "base_agent": base_agent, + "agent_config": agent_config, + "system_instruction": instructions if sent_only_instructions else system_instruction, + "base_environment": base_environment, + "description": description, + "tools": tools, + } + ), + drop_invalid_fields, + MappingProxyType( + { + "id": "name" if id is None else "id", + "system_instruction": "instructions" if sent_only_instructions else "system_instruction", + } + ), + ) + + @client async def acreate( - name: str, + id: str | None = None, base_agent: str | None = None, instructions: str | None = None, base_environment: InteractionEnvironment | None = None, @@ -104,6 +216,12 @@ async def acreate( extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, + agent_config: Mapping[str, object] | None = None, + system_instruction: str | None = None, + description: str | None = None, + tools: Sequence[Mapping[str, object]] | None = None, **kwargs, ) -> AgentCreateResponse: """Async: Create a managed agent on the provider side.""" @@ -113,7 +231,7 @@ async def acreate( kwargs["acreate_agent"] = True func: Final = partial( create, - name=name, + id=id, base_agent=base_agent, instructions=instructions, base_environment=base_environment, @@ -121,6 +239,11 @@ async def acreate( extra_headers=extra_headers, extra_body=extra_body, timeout=timeout, + name=name, + agent_config=agent_config, + system_instruction=system_instruction, + description=description, + tools=tools, **kwargs, ) ctx: Final = contextvars.copy_context() @@ -130,7 +253,7 @@ async def acreate( return init_response except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider or "gemini", original_exception=e, completion_kwargs=local_vars, @@ -140,7 +263,7 @@ async def acreate( @client def create( - name: str, + id: str | None = None, base_agent: str | None = None, instructions: str | None = None, base_environment: InteractionEnvironment | None = None, @@ -148,40 +271,61 @@ def create( extra_headers: dict[str, str] | None = None, extra_body: Mapping[str, object] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, + agent_config: Mapping[str, object] | None = None, + system_instruction: str | None = None, + description: str | None = None, + tools: Sequence[Mapping[str, object]] | None = None, **kwargs, ) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]: """ Sync: Create a managed agent on the provider side. Args: - name: Name for the agent (required). - base_agent: Base agent to derive from (e.g. "waverunner"). - instructions: System instructions for the agent. + id: Identifier for the agent. Exactly one of ``id`` or its alias ``name`` is required. + base_agent: Base agent to derive from (e.g. "antigravity-preview-09-2026"). + instructions: Alias of ``system_instruction``. base_environment: Environment to fork from — an env_id string or a dict like ``{"type": "remote", "sources": [...]}``. custom_llm_provider: Provider to use, e.g. "gemini". extra_headers: Additional HTTP headers. extra_body: Additional request body fields. timeout: Request timeout. + name: Alias of ``id``. + agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}``. + system_instruction: System instruction for the agent. + description: Description of the agent for developers. + tools: Tools available to the agent, e.g. ``[{"type": "code_execution"}]``. **kwargs: Forwarded to GenericLiteLLMParams (api_key, api_base, etc.). """ local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("acreate_agent", False) is True - if base_agent is not None: - kwargs["base_agent"] = base_agent - if instructions is not None: - kwargs["instructions"] = instructions - if base_environment is not None: - kwargs["base_environment"] = base_environment + create_request: Final = _create_request_or_error( + id=id, + name=name, + base_agent=base_agent, + agent_config=agent_config, + system_instruction=system_instruction, + instructions=instructions, + base_environment=base_environment, + description=description, + tools=tools, + drop_invalid_fields=drops_unsupported_params(kwargs.get("drop_params")), + ) + if isinstance(create_request, _InvalidAgentRequest): + raise litellm.BadRequestError( + message=create_request.message, model=id or name or "", llm_provider=custom_llm_provider + ) kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params: Final = GenericLiteLLMParams(**kwargs) - logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "create_agent", {}) + logging_obj: Final = _make_logging_obj(kwargs, create_request.id, custom_llm_provider, "create_agent", {}) config: Final = _get_agents_api_config(custom_llm_provider) return agents_http_handler.create_agent( agents_api_config=config, - name=name, + create_request=create_request, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -191,7 +335,7 @@ def create( ) except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=local_vars, @@ -247,7 +391,7 @@ def list( ) -> AgentListResponse | Coroutine[object, object, AgentListResponse]: """Sync: List all agents on the provider side.""" local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("alist_agents", False) is True kwargs.setdefault("custom_llm_provider", custom_llm_provider) @@ -279,19 +423,22 @@ def list( @client async def aget( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentCreateResponse: - """Async: Get a specific agent by name.""" + """Async: Get a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() try: loop: Final = asyncio.get_event_loop() kwargs["aget_agent"] = True func: Final = partial( get, + id=id, name=name, custom_llm_provider=custom_llm_provider or "gemini", extra_headers=extra_headers, @@ -305,7 +452,7 @@ async def aget( return init_response except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider or "gemini", original_exception=e, completion_kwargs=local_vars, @@ -315,24 +462,27 @@ async def aget( @client def get( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentCreateResponse | Coroutine[object, object, AgentCreateResponse]: - """Sync: Get a specific agent by name.""" + """Sync: Get a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("aget_agent", False) is True + agent_id: Final = _required_agent_id(id, name, custom_llm_provider) kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params: Final = GenericLiteLLMParams(**kwargs) - logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "get_agent", {"name": name}) + logging_obj: Final = _make_logging_obj(kwargs, agent_id, custom_llm_provider, "get_agent", {"name": agent_id}) config: Final = _get_agents_api_config(custom_llm_provider) return agents_http_handler.get_agent( agents_api_config=config, - name=name, + name=agent_id, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -341,7 +491,7 @@ def get( ) except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=local_vars, @@ -356,19 +506,22 @@ def get( @client async def adelete( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentDeleteResult: - """Async: Delete a specific agent by name.""" + """Async: Delete a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() try: loop: Final = asyncio.get_event_loop() kwargs["adelete_agent"] = True func: Final = partial( delete, + id=id, name=name, custom_llm_provider=custom_llm_provider or "gemini", extra_headers=extra_headers, @@ -382,7 +535,7 @@ async def adelete( return init_response except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider or "gemini", original_exception=e, completion_kwargs=local_vars, @@ -392,24 +545,29 @@ async def adelete( @client def delete( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentDeleteResult | Coroutine[object, object, AgentDeleteResult]: - """Sync: Delete a specific agent by name.""" + """Sync: Delete a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("adelete_agent", False) is True + agent_id: Final = _required_agent_id(id, name, custom_llm_provider) kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params: Final = GenericLiteLLMParams(**kwargs) - logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "delete_agent", {"name": name}) + logging_obj: Final = _make_logging_obj( + kwargs, agent_id, custom_llm_provider, "delete_agent", {"name": agent_id} + ) config: Final = _get_agents_api_config(custom_llm_provider) return agents_http_handler.delete_agent( agents_api_config=config, - name=name, + name=agent_id, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -418,7 +576,7 @@ def delete( ) except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=local_vars, @@ -433,19 +591,22 @@ def delete( @client async def alist_versions( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentVersionsResponse: - """Async: List versions of a specific agent.""" + """Async: List versions of a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() try: loop: Final = asyncio.get_event_loop() kwargs["alist_agent_versions"] = True func: Final = partial( list_versions, + id=id, name=name, custom_llm_provider=custom_llm_provider or "gemini", extra_headers=extra_headers, @@ -459,7 +620,7 @@ async def alist_versions( return init_response except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider or "gemini", original_exception=e, completion_kwargs=local_vars, @@ -469,24 +630,29 @@ async def alist_versions( @client def list_versions( - name: str, + id: str | None = None, custom_llm_provider: str | None = None, extra_headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None, + *, + name: str | None = None, **kwargs, ) -> AgentVersionsResponse | Coroutine[object, object, AgentVersionsResponse]: - """Sync: List versions of a specific agent.""" + """Sync: List versions of a specific agent by id (``name`` is an alias of ``id``).""" local_vars: Final = locals() - custom_llm_provider = custom_llm_provider or kwargs.get("custom_llm_provider") or "gemini" + custom_llm_provider = custom_llm_provider or "gemini" try: _is_async: Final = kwargs.pop("alist_agent_versions", False) is True + agent_id: Final = _required_agent_id(id, name, custom_llm_provider) kwargs.setdefault("custom_llm_provider", custom_llm_provider) litellm_params: Final = GenericLiteLLMParams(**kwargs) - logging_obj: Final = _make_logging_obj(kwargs, name, custom_llm_provider, "list_agent_versions", {"name": name}) + logging_obj: Final = _make_logging_obj( + kwargs, agent_id, custom_llm_provider, "list_agent_versions", {"name": agent_id} + ) config: Final = _get_agents_api_config(custom_llm_provider) return agents_http_handler.list_agent_versions( agents_api_config=config, - name=name, + name=agent_id, litellm_params=litellm_params, logging_obj=logging_obj, extra_headers=extra_headers, @@ -495,7 +661,7 @@ def list_versions( ) except Exception as e: raise litellm.exception_type( - model=name, + model=id or name or "", custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=local_vars, diff --git a/litellm/interactions/main.py b/litellm/interactions/main.py index 498e6e20aec..acd43c9b830 100644 --- a/litellm/interactions/main.py +++ b/litellm/interactions/main.py @@ -49,6 +49,7 @@ from litellm.interactions.utils import ( InteractionsAPIRequestUtils, get_provider_interactions_api_config, ) +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.interactions import ( CancelInteractionResult, @@ -100,6 +101,8 @@ async def acreate( timeout: float | httpx.Timeout | None = None, # LiteLLM params custom_llm_provider: str | None = None, + *, + agent_config: Mapping[str, object] | None = None, **kwargs, ) -> InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]: """ @@ -121,6 +124,7 @@ async def acreate( string, or a config object such as ``{"type": "remote", "sources": [...]}`` / ``{"type": "remote", "network": {...}}`` + agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}`` response_modalities: Requested response modalities (TEXT, IMAGE, AUDIO) response_format: JSON schema for response format response_mime_type: MIME type of the response @@ -155,6 +159,7 @@ async def acreate( store=store, background=background, environment=environment, + agent_config=agent_config, response_modalities=response_modalities, response_format=response_format, response_mime_type=response_mime_type, @@ -225,6 +230,8 @@ def create( timeout: float | httpx.Timeout | None = None, # LiteLLM params custom_llm_provider: str | None = None, + *, + agent_config: Mapping[str, object] | None = None, **kwargs, ) -> ( InteractionsAPIResponse @@ -250,6 +257,7 @@ def create( string, or a config object such as ``{"type": "remote", "sources": [...]}`` / ``{"type": "remote", "network": {...}}`` + agent_config: Agent runtime config, e.g. ``{"type": "antigravity", "model": "gemini-3.8-flash"}`` response_modalities: Requested response modalities (TEXT, IMAGE, AUDIO) response_format: JSON schema for response format response_mime_type: MIME type of the response @@ -300,6 +308,15 @@ def create( # Check if this is a bridge provider (litellm_responses) - similar to responses API # Either provider is explicitly "litellm_responses" or no config found (bridge to responses) if custom_llm_provider == "litellm_responses" or interactions_api_config is None: + if optional_params.get("agent_config") is not None and not drops_unsupported_params( + kwargs.get("drop_params") + ): + raise litellm.UnsupportedParamsError( + message=f"'agent_config' is not supported for {custom_llm_provider} interactions. " + "Set drop_params to drop it.", + llm_provider=custom_llm_provider, + model=model, + ) # Bridge to litellm.responses() for non-native providers from litellm.interactions.litellm_responses_transformation.handler import ( LiteLLMResponsesInteractionsHandler, diff --git a/litellm/interactions/streaming_iterator.py b/litellm/interactions/streaming_iterator.py index f4ad57bf7b4..c7faa3fc510 100644 --- a/litellm/interactions/streaming_iterator.py +++ b/litellm/interactions/streaming_iterator.py @@ -21,6 +21,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.interactions import ( InteractionsAPIStreamingResponse, ) @@ -145,7 +146,7 @@ class InteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator): litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, ) - self.stream_iterator = response.aiter_lines() + self.stream_iterator = aiter_sse_lines(response.aiter_text()) def __aiter__(self): return self @@ -212,7 +213,7 @@ class SyncInteractionsAPIStreamingIterator(BaseInteractionsAPIStreamingIterator) litellm_metadata=litellm_metadata, custom_llm_provider=custom_llm_provider, ) - self.stream_iterator = response.iter_lines() + self.stream_iterator = iter_sse_lines(response.iter_text()) def __iter__(self): return self diff --git a/litellm/litellm_core_utils/README.md b/litellm/litellm_core_utils/README.md index a5f5e8326b9..331f4adccbb 100644 --- a/litellm/litellm_core_utils/README.md +++ b/litellm/litellm_core_utils/README.md @@ -11,4 +11,4 @@ Core files: - `duration_parser.py`: code for parsing durations - e.g. "1d", "1mo", "10s" - `api_route_to_call_types.py`: mapping of API routes to their corresponding CallTypes (e.g., `/chat/completions` -> [acompletion, completion]) -Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Missing native bindings or unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer. +Tokenizer factories return Python tokenizer objects by default. Set `LITELLM_RUST=1` or call `litellm.rust(True)` before constructing tokenizers to select the Rust backend through `Route.TOKENIZER` in the Rust catalog. Unsupported native features fall back to Python. Existing tokenizer objects keep their selected backend. Rust-backed tokenizer objects carry the read-only `tiktoken.Encoding` / `tokenizers.Tokenizer` surface and are immutable: `enable_padding`, `enable_truncation` and `add_tokens` stay on the Python tokenizer. diff --git a/litellm/litellm_core_utils/audio_utils/utils.py b/litellm/litellm_core_utils/audio_utils/utils.py index dab3e48f91a..640559ac7e2 100644 --- a/litellm/litellm_core_utils/audio_utils/utils.py +++ b/litellm/litellm_core_utils/audio_utils/utils.py @@ -2,8 +2,12 @@ Utils used for litellm.transcription() and litellm.atranscription() """ +import base64 +import binascii import hashlib +import math import os +from collections.abc import Sequence from dataclasses import dataclass from typing import Final @@ -331,6 +335,45 @@ def calculate_request_duration(file: FileTypes) -> float | None: return None +DEFAULT_AUDIO_TOKEN_ESTIMATE: Final = 300 +AUDIO_TOKENS_PER_SECOND: Final = 10 +AUDIO_BYTES_PER_TOKEN: Final = 1600 +_MAX_AUDIO_ESTIMATE_BYTES: Final = 16 * 1024 * 1024 + + +def _audio_token_rate(model: str) -> int | None: + if model.startswith( + ("gemini-", "gemini/gemini-", "vertex_ai/gemini-", "vertex_ai_beta/gemini-", "openrouter/google/gemini-") + ): + family: Final = model.rsplit("/", 1)[-1] + return 25 if family.startswith(("gemini-3-", "gemini-3.")) else 32 + if model.startswith(("gpt-", "openai/gpt-", "azure/gpt-", "openrouter/openai/gpt-")): + return AUDIO_TOKENS_PER_SECOND + return None + + +def estimate_audio_tokens(data: str | None, models: Sequence[str] = (), fallback_model: str | None = None) -> int: + """Estimate audio tokens from container duration within a bounded decoding budget.""" + if not data: + return DEFAULT_AUDIO_TOKEN_ESTIMATE + fallback_estimate: Final = max(len(data) * 3 // 4 // AUDIO_BYTES_PER_TOKEN, DEFAULT_AUDIO_TOKEN_ESTIMATE) + padding: Final = 2 if data.endswith("==") else 1 if data.endswith("=") else 0 + if len(data) * 3 // 4 - padding > _MAX_AUDIO_ESTIMATE_BYTES: + return fallback_estimate + try: + duration: Final = calculate_request_duration(base64.b64decode(data, validate=True)) + except (ValueError, binascii.Error): + return fallback_estimate + if duration is None or not math.isfinite(duration) or duration <= 0: + return fallback_estimate + fallback_rate: Final = _audio_token_rate(fallback_model) if fallback_model is not None else None + rate: Final = max( + (_audio_token_rate(model) or fallback_rate or AUDIO_TOKENS_PER_SECOND for model in models), + default=fallback_rate or AUDIO_TOKENS_PER_SECOND, + ) + return max(1, math.ceil(duration * rate)) + + DEFAULT_SPEECH_MEDIA_TYPE: Final = "audio/mpeg" diff --git a/litellm/router_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py similarity index 100% rename from litellm/router_utils/batch_utils.py rename to litellm/litellm_core_utils/batch_utils.py diff --git a/litellm/litellm_core_utils/cli_token_utils.py b/litellm/litellm_core_utils/cli_token_utils.py index ffe2c4433a0..37ed4efa0ca 100644 --- a/litellm/litellm_core_utils/cli_token_utils.py +++ b/litellm/litellm_core_utils/cli_token_utils.py @@ -353,9 +353,9 @@ def is_cli_token_fresh( buffer_hours: float = CLI_TOKEN_FRESHNESS_BUFFER_SECONDS / 3600, ) -> bool: """Check whether a cached CLI token is still within its expiration window. - Used by `lite auth print-token` to fail fast, without a network round trip, - once the cached token is past `LITELLM_CLI_JWT_EXPIRATION_HOURS`. A `--pkce` - credential carries its own `expires_at`, which is authoritative when present.""" + Used by `lite auth print-token` to skip asking the proxy while the cached token + is younger than `LITELLM_CLI_JWT_EXPIRATION_HOURS`. A `--pkce` credential + carries its own `expires_at`, which is authoritative when present.""" from litellm.constants import CLI_JWT_EXPIRATION_HOURS expires_at: Final = ( diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 00c6ce97e78..55025508e40 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -266,6 +266,7 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { "tool_use": "tool_calls", "refusal": "content_filter", "compaction": "length", + "model_context_window_exceeded": "length", # Cohere "COMPLETE": "stop", "ERROR_TOXIC": "content_filter", diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index ba964d200ec..8178b7e98b8 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -75,6 +75,8 @@ OPTIONAL_KWARGS_KEYS: Final = ( "rpm", "itpm", "otpm", + "tpd", + "rpd", "use_xai_oauth", "chatgpt_auth_file", "github_copilot_auth_type", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 19c1d5c414c..d9933087334 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -437,7 +437,6 @@ def get_llm_provider( model in litellm.open_ai_chat_completion_models or "ft:gpt-3.5-turbo" in model or "ft:gpt-4" in model # catches ft:gpt-4-0613, ft:gpt-4o - or model in litellm.openai_image_generation_models or model.startswith("gpt-image") or model in litellm.openai_video_generation_models ): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 48e8a44b9a4..8b66c27b930 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -842,6 +842,7 @@ class Logging(LiteLLMLoggingBaseClass): self._async_success_enqueued: bool = False self._sync_success_scheduled: bool = False self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None + self._stream_failure_claimed: bool = False self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging": @@ -2561,6 +2562,12 @@ class Logging(LiteLLMLoggingBaseClass): }, ) + def claim_stream_failure_report(self) -> bool: + if self._stream_failure_claimed: + return False + self._stream_failure_claimed = True + return True + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" usage: Final = getattr(assembled, "usage", None) @@ -2751,7 +2758,7 @@ class Logging(LiteLLMLoggingBaseClass): if entry is None or not is_str_keyed_dict(litellm_params): return metadata: Final[object] = litellm_params.get("metadata") - bucket: Final = metadata if is_str_keyed_dict(metadata) else {} + bucket: Final = {**metadata} if is_str_keyed_dict(metadata) else {} litellm_params["metadata"] = bucket with_guardrail_entry(bucket, entry) @@ -3474,7 +3481,7 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, print_verbose=print_verbose, - level=LogfireLevel.INFO.value, + level=LogfireLevel.INFO, ) if callback == "lunary" and lunaryLogger is not None: @@ -4435,7 +4442,7 @@ class Logging(LiteLLMLoggingBaseClass): response_obj=result, start_time=start_time, end_time=end_time, - level=LogfireLevel.ERROR.value, + level=LogfireLevel.ERROR, print_verbose=print_verbose, ) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 3b49f239afe..4b0b875af6d 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field, replace from datetime import datetime, timezone, tzinfo from types import MappingProxyType -from typing import Final, Literal, TypedDict, cast +from typing import Final, Literal, Protocol, TypedDict, cast from zoneinfo import ZoneInfo, ZoneInfoNotFoundError from typing_extensions import ReadOnly @@ -2085,6 +2085,41 @@ def calculate_image_response_web_search_cost( ) +class _ModelPricedImageCostCalculator(Protocol): + def __call__(self, model: str, image_response: ImageResponse, model_info: ModelInfo | None = None) -> float: ... + + +def _model_priced_image_cost_calculator( + custom_llm_provider: str | None, model: str +) -> _ModelPricedImageCostCalculator | None: + match custom_llm_provider: + case litellm.LlmProviders.RECRAFT.value: + from litellm.llms.recraft.cost_calculator import cost_calculator as recraft_image_cost_calculator + + return recraft_image_cost_calculator + case litellm.LlmProviders.AIML.value: + from litellm.llms.aiml.image_generation.cost_calculator import cost_calculator as aiml_image_cost_calculator + + return aiml_image_cost_calculator + case litellm.LlmProviders.COMETAPI.value: + from litellm.llms.cometapi.image_generation.cost_calculator import ( + cost_calculator as cometapi_image_cost_calculator, + ) + + return cometapi_image_cost_calculator + case litellm.LlmProviders.RUNWAYML.value: + from litellm.llms.runwayml.cost_calculator import cost_calculator as runwayml_image_cost_calculator + + return runwayml_image_cost_calculator + case litellm.LlmProviders.TENCENT.value: + from litellm.llms.tencent.cost_calculator import image_cost_calculator as tencent_image_cost_calculator + from litellm.llms.tencent.image_generation.transformation import is_tokenhub_image_model + + return tencent_image_cost_calculator if is_tokenhub_image_model(model) else None + case _: + return None + + class CostCalculatorUtils: @staticmethod def call_type_has_image_response(call_type: str) -> bool: @@ -2131,9 +2166,6 @@ class CostCalculatorUtils: from litellm.llms.gemini.image_generation.cost_calculator import ( cost_calculator as gemini_image_cost_calculator, ) - from litellm.llms.recraft.cost_calculator import ( - cost_calculator as recraft_image_cost_calculator, - ) from litellm.llms.vertex_ai.image_generation.cost_calculator import ( cost_calculator as vertex_ai_image_cost_calculator, ) @@ -2146,7 +2178,10 @@ class CostCalculatorUtils: ) resolved_n: Final = n if n is not None else (len(completion_response.data) if completion_response.data else 0) pricing: Final = deployment_pricing(model_info) + model_priced_calculator: Final = _model_priced_image_cost_calculator(custom_llm_provider, model) + if model_priced_calculator is not None: + return model_priced_calculator(model=model, image_response=completion_response, model_info=pricing) if custom_llm_provider == litellm.LlmProviders.VERTEX_AI.value: if isinstance(completion_response, ImageResponse): return vertex_ai_image_cost_calculator( @@ -2164,36 +2199,6 @@ class CostCalculatorUtils: optional_params=optional_params, ) raise TypeError("completion_response must be of type ImageResponse for bedrock image cost calculation") - elif custom_llm_provider == litellm.LlmProviders.RECRAFT.value: - from litellm.llms.recraft.cost_calculator import ( - cost_calculator as recraft_image_cost_calculator, - ) - - return recraft_image_cost_calculator( - model=model, - image_response=completion_response, - model_info=pricing, - ) - elif custom_llm_provider == litellm.LlmProviders.AIML.value: - from litellm.llms.aiml.image_generation.cost_calculator import ( - cost_calculator as aiml_image_cost_calculator, - ) - - return aiml_image_cost_calculator( - model=model, - image_response=completion_response, - model_info=pricing, - ) - elif custom_llm_provider == litellm.LlmProviders.COMETAPI.value: - from litellm.llms.cometapi.image_generation.cost_calculator import ( - cost_calculator as cometapi_image_cost_calculator, - ) - - return cometapi_image_cost_calculator( - model=model, - image_response=completion_response, - model_info=pricing, - ) elif custom_llm_provider == litellm.LlmProviders.GEMINI.value: if call_type in ( CallTypes.image_edit.value, @@ -2237,16 +2242,6 @@ class CostCalculatorUtils: optional_params=optional_params, model_info=pricing, ) - elif custom_llm_provider == litellm.LlmProviders.RUNWAYML.value: - from litellm.llms.runwayml.cost_calculator import ( - cost_calculator as runwayml_image_cost_calculator, - ) - - return runwayml_image_cost_calculator( - model=model, - image_response=completion_response, - model_info=pricing, - ) elif ( custom_llm_provider == litellm.LlmProviders.OPENAI.value or custom_llm_provider == litellm.LlmProviders.AZURE.value diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 3c06c522dd6..101f069c0f8 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -9,6 +9,7 @@ import json import mimetypes import re from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass from graphlib import CycleError, TopologicalSorter from itertools import chain, groupby, islice from os import PathLike @@ -20,8 +21,8 @@ from pydantic import TypeAdapter import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.batch_utils import InMemoryFile from litellm.litellm_core_utils.core_helpers import as_str_mapping -from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -236,6 +237,160 @@ _IMAGE_CONTENT_PART_TYPES: Final = frozenset({"image_url", "input_image", "image _IMAGE_SCAN_MAX_DEPTH: Final = 4 +@dataclass(frozen=True, slots=True) +class _TPMMediaProjection: + content: object + media: tuple[Mapping[str, object], ...] = () + + +def _tpm_image_part(url: str, detail: object = "auto") -> Mapping[str, object]: + return {"type": "image_url", "image_url": {"url": url, "detail": detail}} + + +def _tpm_blank_fields(fields: Mapping[str, object], names: tuple[str, ...]) -> Mapping[str, object]: + return {key: "" if key in names and isinstance(value, str) else value for key, value in fields.items()} + + +def _tpm_file_media(part: Mapping[str, object]) -> _TPMMediaProjection: + from litellm.litellm_core_utils.prompt_templates.factory import parse_mime_type + + nested: Final = part.get("type") == "file" + fields: Final = as_str_mapping(part.get("file")) if nested else part + if fields is None: + return _TPMMediaProjection(part) + raw_data: Final = fields.get("file_data") + data: Final = raw_data if isinstance(raw_data, str) else "" + declared: Final = next( + ( + value + for key in ("format", "mime_type", "content_type") + if isinstance(value := fields.get(key), str) and value + ), + "", + ) + mime: Final = (declared or parse_mime_type(data[:256]) or "").split(";", 1)[0].strip().lower() + if not mime.startswith(("image/", "audio/")): + return _TPMMediaProjection(part) + raw_url: Final = fields.get("file_url") or fields.get("file_id") + url: Final = data or (raw_url if isinstance(raw_url, str) else "") + header_end: Final = data.find(",", 0, 256) + encoded: Final = ( + data[header_end + 1 :] + if header_end >= 0 and data[:header_end].endswith(";base64") + else "" + if data.startswith(("data:", "https://", "http://", "gs://")) + else data + ) + media: Final = ( + _tpm_image_part( + url + if not data or url.startswith(("data:", "https://", "http://", "gs://")) + else f"data:{mime};base64,{url}", + fields.get("detail", "auto"), + ) + if mime.startswith("image/") + else {"type": "input_audio", "input_audio": {"data": encoded, "format": mime.removeprefix("audio/")}} + ) + projected: Final = _tpm_blank_fields(fields, ("file_data", "file_url")) + return _TPMMediaProjection({**part, "file": projected} if nested else projected, (media,)) + + +def _tpm_media_leaf(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + if part is None: + return _TPMMediaProjection(value) + kind: Final = part.get("type") + if kind in ("file", "input_file"): + return _tpm_file_media(part) + if kind == "input_audio": + audio: Final = as_str_mapping(part.get("input_audio")) + return _TPMMediaProjection( + {**part, "input_audio": _tpm_blank_fields(audio, ("data",))} if audio is not None else value, + ({"type": "input_audio", "input_audio": part.get("input_audio")},), + ) + if kind in ("image_url", "input_image", "image", "computer_screenshot"): + image: Final = as_str_mapping(part.get("image_url")) + source: Final = as_str_mapping(part.get("source")) + canonical: Final = _tpm_image_part( + _image_part_url({**part, "type": "input_image"} if kind == "computer_screenshot" else part) or "", + (image or part).get("detail", part.get("detail", "auto")), + ) + projected: Final = ( + {**part, "image_url": _tpm_blank_fields(image, ("url",))} + if image is not None and kind in ("image_url", "input_image", "computer_screenshot") + else {**part, "source": _tpm_blank_fields(source, ("data", "url"))} + if source is not None and kind == "image" and source.get("type") in ("base64", "url") + else _tpm_blank_fields(part, ("image_url", "url")) + ) + return _TPMMediaProjection(projected, (canonical,)) + return _TPMMediaProjection(value) + + +def _tpm_media_list(value: object, split: Callable[[object], _TPMMediaProjection]) -> _TPMMediaProjection: + if not isinstance(value, (list, tuple)): + return _TPMMediaProjection(value) + parts: Final = cast(Sequence[object], value) # cast-ok: finite JSON content arrays + projected: Final = tuple(map(split, parts)) + media: Final = tuple(chain.from_iterable(part.media for part in projected)) + content: Final = tuple(part.content for part in projected) + return _TPMMediaProjection(parts if not media else content if isinstance(value, tuple) else list(content), media) + + +def _tpm_replace_field(part: Mapping[str, object], key: str, projection: _TPMMediaProjection) -> _TPMMediaProjection: + return _TPMMediaProjection({**part, key: projection.content} if projection.media else part, projection.media) + + +def _tpm_document_media(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + source: Final = as_str_mapping(part.get("source")) if part is not None else None + if part is None or part.get("type") != "document" or source is None or source.get("type") != "content": + return _tpm_media_leaf(value) + content: Final = _tpm_media_list(source.get("content"), _tpm_media_leaf) + return _tpm_replace_field(part, "source", _tpm_replace_field(source, "content", content)) + + +def _tpm_tool_result_media(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + if part is None or part.get("type") != "tool_result": + return _tpm_document_media(value) + return _tpm_replace_field(part, "content", _tpm_media_list(part.get("content"), _tpm_document_media)) + + +def _tpm_code_output_media(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + url: Final = part.get("url") if part is not None else None + if part is None or part.get("type") != "image" or not isinstance(url, str) or not url: + return _TPMMediaProjection(value) + return _TPMMediaProjection(_tpm_blank_fields(part, ("url",)), (_tpm_image_part(url),)) + + +def _tpm_input_media(value: object) -> _TPMMediaProjection: + part: Final = as_str_mapping(value) + if part is None: + return _TPMMediaProjection(value) + kind: Final = part.get("type") + if kind in ("function_call_output", "custom_tool_call_output"): + return _tpm_replace_field(part, "output", _tpm_media_list(part.get("output"), _tpm_tool_result_media)) + if kind == "computer_call_output": + return _tpm_replace_field(part, "output", _tpm_media_leaf(part.get("output"))) + if kind == "code_interpreter_call": + return _tpm_replace_field(part, "outputs", _tpm_media_list(part.get("outputs"), _tpm_code_output_media)) + result: Final = part.get("result") + if kind == "image_generation_call" and isinstance(result, str) and result: + return _TPMMediaProjection( + _tpm_blank_fields(part, ("result",)), + (_tpm_image_part("data:application/octet-stream;base64," + result),), + ) + if kind == "message" or (kind is None and isinstance(part.get("role"), str)): + return _tpm_replace_field(part, "content", _tpm_media_list(part.get("content"), _tpm_tool_result_media)) + return _tpm_tool_result_media(value) + + +def split_media_for_tpm(value: object) -> tuple[object, tuple[Mapping[str, object], ...]]: + projection: Final = _tpm_media_list(value, _tpm_input_media) + return projection.content, projection.media + + def _content_parts_contain_image(parts: Sequence[object]) -> bool: """Depth-bounded frontier walk over nested content lists, iterative because the repo bans recursion; an Anthropic tool_result nests its image parts exactly one level down.""" diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 63242a580e7..24c401281ef 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -34,7 +34,7 @@ def safe_json_structure( def _transform(key: str | None, value: str) -> str: return value if value_transform is None else value_transform(key, value) - def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> Any: + def _serialize(obj: object, seen: set[int], depth: int, key: str | None = None) -> object: # Check for maximum depth. if depth > max_depth: return "MaxDepthExceeded" @@ -48,32 +48,32 @@ def safe_json_structure( if id(obj) in seen: return "CircularReference Detected" seen.add(id(obj)) - result: dict | list | tuple | set | str if isinstance(obj, dict): - result = {} - for k, v in obj.items(): - if isinstance(k, (str)): - clean_k = k.replace("\x00", "") if "\x00" in k else k - result[clean_k] = _serialize(v, seen, depth + 1, clean_k) + mapping_result: Final = { + (clean_k := k.replace("\x00", "") if "\x00" in k else k): _serialize(v, seen, depth + 1, clean_k) + for k, v in obj.items() + if isinstance(k, str) + } seen.remove(id(obj)) - return result + return mapping_result elif isinstance(obj, list): - result = [_serialize(item, seen, depth + 1, key) for item in obj] + list_result: Final = [_serialize(item, seen, depth + 1, key) for item in obj] seen.remove(id(obj)) - return result + return list_result elif isinstance(obj, tuple): - result = tuple(_serialize(item, seen, depth + 1, key) for item in obj) + tuple_result: Final = tuple(_serialize(item, seen, depth + 1, key) for item in obj) seen.remove(id(obj)) - return result + return tuple_result elif isinstance(obj, set): - result = sorted([_serialize(item, seen, depth + 1, key) for item in obj]) + sortable_values: Final[list[Any]] = [_serialize(item, seen, depth + 1, key) for item in obj] + set_result: Final = sorted(sortable_values) seen.remove(id(obj)) - return result + return set_result elif isinstance(obj, BaseModel): dumped: Final = obj.model_dump() - result = _serialize(dumped, seen, depth + 1, key) + model_result: Final = _serialize(dumped, seen, depth + 1, key) seen.remove(id(obj)) - return result + return model_result else: # Fall back to string conversion for non-serializable objects. try: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b64e146ca08..2348d8e57d6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1263,7 +1263,7 @@ class CustomStreamWrapper: ) if _is_delta_empty: - model_response.choices[0].delta = Delta(content=None) # ensure empty delta chunk returned + model_response.choices[0].delta = Delta(content=self._close_open_think_block()) # get any function call arguments model_response.choices[0].finish_reason = map_finish_reason( finish_reason=self.received_finish_reason @@ -1292,21 +1292,26 @@ class CustomStreamWrapper: if self.merge_reasoning_content_in_choices is not True: return delta: Final = model_response.choices[0].delta - reasoning_content: Final = getattr(delta, "reasoning_content", None) + reasoning_content: Final = getattr(delta, "reasoning_content", None) or "" content: Final = delta.content or "" - if reasoning_content: - self._emitted_disqualifying_content = True - closes_thinking: Final = bool(content) and not self.sent_last_thinking_block - opening: Final = "" if self.sent_first_thinking_block else "" - delta.content = opening + reasoning_content + ("" + content if closes_thinking else "") - self.sent_first_thinking_block = True - self.sent_last_thinking_block = self.sent_last_thinking_block or closes_thinking - elif self.sent_first_thinking_block and not self.sent_last_thinking_block and content: - delta.content = "" + content - self.sent_last_thinking_block = True - + ends_thinking: Final = bool(content or delta.tool_calls or delta.function_call) + inside_think_block: Final = self.sent_first_thinking_block and not self.sent_last_thinking_block + opening: Final = "" if reasoning_content and not inside_think_block else "" + closing: Final = "" if ends_thinking and (reasoning_content or inside_think_block) else "" if hasattr(delta, "reasoning_content"): del delta.reasoning_content + if not reasoning_content and not closing: + return + delta.content = opening + reasoning_content + closing + content + self.sent_first_thinking_block = True + self.sent_last_thinking_block = bool(closing) + self._emitted_disqualifying_content = self._emitted_disqualifying_content or bool(reasoning_content) + + def _close_open_think_block(self) -> str | None: + if not self.sent_first_thinking_block or self.sent_last_thinking_block: + return None + self.sent_last_thinking_block = True + return "" def _dispatch_provider_chunk( self, @@ -1557,6 +1562,12 @@ class CustomStreamWrapper: else: # openai / azure chat model if isinstance(chunk, ChatCompletionChunk): self._openai_handler_receives_chat_chunks = True + if any(str(choice.finish_reason) == "abort" for choice in chunk.choices or ()): + raise litellm.InternalServerError( + message="Upstream generation aborted (finish_reason=abort)", + model=chunk.model, + llm_provider=self.custom_llm_provider, + ) if self.custom_llm_provider in [ LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value, @@ -1926,8 +1937,20 @@ class CustomStreamWrapper: model_response.choices[0].finish_reason == "stop" and self.tool_call ): # don't overwrite for other - potential error finish reasons model_response.choices[0].finish_reason = "tool_calls" + close_tag: Final = self._close_open_think_block() + if close_tag is not None: + model_response.choices[0].delta.content = close_tag return model_response + def _final_chunk_for_logging(self, processed_chunk: "ModelResponseStream") -> "ModelResponseStream | None": + if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: + return processed_chunk + if processed_chunk.choices[0].delta.content is None: + return None + return processed_chunk.model_copy( + update={"choices": [processed_chunk.choices[0].model_copy(update={"finish_reason": None})]} + ) + def _record_usage_only_chunk(self, model_response: "ModelResponseStream") -> None: """ Keep provider usage-only chunks (e.g. OpenRouter's post-finish chunk, which carries a @@ -2167,10 +2190,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() - # The logged response is built from self.chunks; keep a finish_reason the provider sent on its - # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. - if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: - self.chunks.append(processed_chunk) + final_chunk_for_logging: Final = self._final_chunk_for_logging(processed_chunk) + if final_chunk_for_logging is not None: + self.chunks.append(final_chunk_for_logging) if self.stream_options is None: self._set_final_chunk_usage(processed_chunk) ## LOGGING @@ -2400,10 +2422,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() - # The logged response is built from self.chunks; keep a finish_reason the provider sent on its - # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. - if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: - self.chunks.append(processed_chunk) + final_chunk_for_logging: Final = self._final_chunk_for_logging(processed_chunk) + if final_chunk_for_logging is not None: + self.chunks.append(final_chunk_for_logging) if self.stream_options is None: self._set_final_chunk_usage(processed_chunk) # see sync __next__'s sibling branch: deliberately do NOT restore diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 6d50e4b065e..f42e4eae08f 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -3,9 +3,12 @@ import base64 import io import math +import re import struct +from bisect import bisect_right from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from itertools import accumulate +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, cast import anyio @@ -35,6 +38,7 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( ChatTemplateToolSection, chat_template_tool_section, @@ -188,22 +192,23 @@ def resize_image_high_res( # Test the function with the given example def calculate_tiles_needed( - resized_width, - resized_height, - tile_width=MAX_TILE_WIDTH, - tile_height=MAX_TILE_HEIGHT, -): + resized_width: int, + resized_height: int, + tile_width: int = MAX_TILE_WIDTH, + tile_height: int = MAX_TILE_HEIGHT, +) -> int: tiles_across: Final = (resized_width + tile_width - 1) // tile_width tiles_down: Final = (resized_height + tile_height - 1) // tile_height total_tiles: Final = tiles_across * tiles_down return total_tiles -def high_detail_image_token_upper_bound(base_tokens: int = 85) -> int: +def high_detail_image_token_upper_bound(base_tokens: int = 85, tile_tokens: int | None = None) -> int: largest_tile_count: Final = calculate_tiles_needed( MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES, MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES ) - return base_tokens + (base_tokens * 2) * largest_tile_count + per_tile: Final = base_tokens * 2 if tile_tokens is None else tile_tokens + return base_tokens + per_tile * largest_tile_count def _unpack_ints(fmt: str, buffer: bytes) -> tuple[int, ...]: @@ -381,6 +386,219 @@ def calculate_img_tokens( return total_tokens +_IMAGE_HEADER_BASE64_CHARS: Final = 65_536 +_MAX_IMAGE_PATCHES: Final = 30_000 +_ANTHROPIC_IMAGE_LIMITS: Final = MappingProxyType( + {"claude-standard": (ANTHROPIC_IMAGE_MAX_LONG_EDGE_PX, 1568), "claude-high": (2576, 4784)} +) +_TILE_IMAGE_TOKEN_COSTS: Final = MappingProxyType( + { + "gpt-4o": (85, 170), + "gpt-4.1": (85, 170), + "gpt-4o-mini": (2833, 5667), + "gpt-5": (70, 140), + "gpt-5.1": (70, 140), + "o1": (75, 150), + "o1-pro": (75, 150), + "o3": (75, 150), + } +) +_PATCH_IMAGE_MULTIPLIERS: Final = MappingProxyType( + { + "gpt-6-astra": 1.2, + "gpt-5.6-sol": 1.2, + "gpt-5.6-terra": 1.2, + "gpt-5.6-luna": 1.2, + "gpt-5.5": 1.2, + "gpt-5.4": 1.2, + "gpt-5.4-mini": 1.2, + "gpt-5.4-nano": 1.2, + "gpt-5.2": 1.2, + "gpt-4.1-mini": 1.62, + } +) + + +def _inline_image_dimensions(data: str) -> tuple[int, int] | None: + separator: Final = data.find(",", 0, 256) + if not data.startswith("data:") or separator < 0 or not data[:separator].endswith(";base64"): + return None + encoded_header: Final = data[separator + 1 : separator + 1 + _IMAGE_HEADER_BASE64_CHARS] + try: + header: Final = base64.b64decode(encoded_header, validate=True) + image_type: Final = get_image_type(header) + if image_type not in ("png", "gif", "jpeg", "webp"): + return None + if image_type == "webp" and header[12:16] not in (b"VP8X", b"VP8 ", b"VP8L"): + return None + dimensions: Final = _image_dimensions_from_bytes(header) + except (ValueError, TypeError, struct.error): + return None + width, height = dimensions + return (width, height) if width > 0 and height > 0 else None + + +def _patch_image_limits(model: str, detail: str) -> tuple[int, int | None] | None: + if model in ("gpt-5.2", "gpt-4.1-mini"): + return None if detail == "original" else (2048, 6144) + original_default: Final = model in ("gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.5") + effective_detail: Final = ("original" if original_default else "high") if detail == "auto" else detail + if effective_detail == "low": + return (2048, 6144) if model in ("gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano") else (512, None) + if effective_detail == "high": + return (65_535 if model == "gpt-6-astra" else 2048, 2500) + return ( + (65_535, None) if model in ("gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna") else (6000, 10_000) + ) + + +def _image_patch_count(width: int, height: int, max_dimension: int, budget: int | None) -> int: + dimension_scale: Final = min(1.0, max_dimension / max(width, height)) + fitted_width: Final = max(1, round(width * dimension_scale)) + fitted_height: Final = max(1, round(height * dimension_scale)) + patches: Final = math.ceil(fitted_width / 32) * math.ceil(fitted_height / 32) + if budget is None or patches <= budget: + return patches + shrink: Final = math.sqrt(32**2 * budget / (fitted_width * fitted_height)) + width_patches: Final = fitted_width * shrink / 32 + height_patches: Final = fitted_height * shrink / 32 + aligned_shrink: Final = shrink * min( + math.floor(width_patches) / width_patches, math.floor(height_patches) / height_patches + ) + resized_width: Final = max(1, math.floor(fitted_width * aligned_shrink)) + resized_height: Final = max(1, math.floor(fitted_height * aligned_shrink)) + return math.ceil(resized_width / 32) * math.ceil(resized_height / 32) + + +def _image_policy_model_name(model: str) -> str | None: + provider, separator, provider_model = model.partition("/") + if not separator: + return model + if provider == "openai": + return provider_model + if provider == "openrouter": + return provider_model.removeprefix("openai/") if provider_model.startswith("openai/") else None + if provider not in ("azure", "azure_ai"): + return None + catalog_entry: Final = as_str_mapping( + cast( # cast-ok: read-only view of the legacy string-keyed registry; as_str_mapping checks the entry + Mapping[str, object], litellm.model_cost + ).get(model) + ) + if catalog_entry is not None and catalog_entry.get("litellm_provider") in ("azure", "azure_ai"): + return provider_model.rsplit("/", 1)[-1] + return provider_model.removeprefix("responses/").removeprefix("o_series/").removeprefix("openai/") + + +def _anthropic_image_family(model: str) -> str | None: + from litellm.utils import get_potential_model_names + + provider, separator, _ = model.partition("/") + if separator and provider not in ( + "anthropic", + "bedrock", + "bedrock_converse", + "bedrock_mantle", + "vertex_ai", + "vertex_ai_beta", + "azure_ai", + "openrouter", + ): + return None + candidates: Final = get_potential_model_names(model, provider if separator else "anthropic") + catalog: Final = cast(Mapping[str, object], litellm.model_cost) # cast-ok: read-only view of the legacy registry + for candidate in dict.fromkeys(candidates.values()): + if not isinstance(candidate, str): + continue + if (entry := as_str_mapping(catalog.get(candidate))) is None or entry.get("supports_vision") is not True: + continue + if ( + version := re.search( + r"(?:^|[./])claude-(?:(?:haiku|sonnet|opus)-)?(\d+)(?:[-.](\d{1,2})(?!\d))?", candidate + ) + ) is None: + continue + major, minor = int(version[1]), int(version[2] or 0) + return "claude-high" if (major, minor) >= (4, 7) else "claude-standard" + return None + + +def _anthropic_image_tokens(dimensions: tuple[int, int] | None, limits: tuple[int, int]) -> int: + max_edge, token_budget = limits + if dimensions is None: + return token_budget + long_edge, short_edge = max(dimensions), min(dimensions) + aspect_ratio: Final = long_edge / short_edge + + def tokens_at_edge(edge: int) -> int: + return math.ceil(edge / 28) * math.ceil(max(round(edge / aspect_ratio), 1) / 28) + + fitted_edge: Final = bisect_right(range(1, min(long_edge, max_edge) + 1), token_budget, key=tokens_at_edge) + return tokens_at_edge(max(fitted_edge, 1)) + + +def _image_model_family(model: str) -> str | None: + provider_model: Final = _image_policy_model_name(model) + if provider_model is None: + return _anthropic_image_family(model) + family: Final = re.sub(r"-\d{4}-\d{2}-\d{2}$", "", provider_model) + canonical: Final = litellm.azure_llms.get(f"azure/{family}", family) if model.startswith("azure/") else family + return ( + canonical + if canonical in _TILE_IMAGE_TOKEN_COSTS or canonical in _PATCH_IMAGE_MULTIPLIERS + else _anthropic_image_family(model) + ) + + +def _image_policy_tokens(family: str | None, dimensions: tuple[int, int] | None, detail: str) -> int: + if family is None: + return DEFAULT_IMAGE_TOKEN_COUNT + anthropic_limits: Final = _ANTHROPIC_IMAGE_LIMITS.get(family) + if anthropic_limits is not None: + return _anthropic_image_tokens(dimensions, anthropic_limits) + tile_costs: Final = _TILE_IMAGE_TOKEN_COSTS.get(family) + multiplier: Final = _PATCH_IMAGE_MULTIPLIERS.get(family) + if tile_costs is not None: + base_tokens, tile_tokens = tile_costs + if detail == "original": + return DEFAULT_IMAGE_TOKEN_COUNT + if detail == "low": + return base_tokens + if dimensions is None: + return high_detail_image_token_upper_bound(base_tokens, tile_tokens) + width, height = dimensions + scale: Final = min(1.0, 2048 / max(width, height), 768 / min(width, height)) + tiles: Final = calculate_tiles_needed(max(1, math.floor(width * scale)), max(1, math.floor(height * scale))) + return base_tokens + tiles * tile_tokens + limits: Final = _patch_image_limits(family, detail) + if limits is None or multiplier is None: + return DEFAULT_IMAGE_TOKEN_COUNT + max_dimension, patch_budget = limits + patches: Final = ( + min(math.ceil(max_dimension / 32) ** 2, patch_budget if patch_budget is not None else _MAX_IMAGE_PATCHES) + if dimensions is None + else _image_patch_count(*dimensions, max_dimension, patch_budget) + ) + return math.ceil(patches * multiplier) + + +def estimate_image_tokens( + model: str | Sequence[str], data: str, detail: str = "auto", *, fallback_model: str | None = None +) -> int: + """Estimate image policies from at most 48 KiB of inline headers, without fetching media.""" + if detail not in ("low", "high", "auto", "original"): + return DEFAULT_IMAGE_TOKEN_COUNT + models: Final = (model,) if isinstance(model, str) else model + fallback_family: Final = _image_model_family(fallback_model) if fallback_model is not None else None + families: Final = tuple( + dict.fromkeys(_image_model_family(candidate) or fallback_family for candidate in models or ("",)) + ) + if all(family is None for family in families): + return DEFAULT_IMAGE_TOKEN_COUNT + dimensions: Final = _inline_image_dimensions(data) + return max(_image_policy_tokens(family, dimensions, detail) for family in families) + + TokenCounterFunction = Callable[[str], int] """ Type for a function that counts tokens in a string. diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index 74f782b56e7..9253ca389ee 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -82,4 +82,6 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): model_response.model = _json_response.get("model") model_response.object = _json_response.get("object") model_response.system_fingerprint = _json_response.get("system_fingerprint") + if "service_tier" in _json_response: + return model_response.model_copy(update={"service_tier": _json_response.get("service_tier")}) return model_response diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 13ceea3e31b..8a82ed98991 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -38,6 +38,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.llms.anthropic import ( ContentBlockDelta, ContentBlockStart, @@ -233,7 +234,7 @@ async def make_call( raise AnthropicError(status_code=500, message=str(e)) completion_stream: Final = ModelResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, json_mode=json_mode, speed=speed, @@ -304,7 +305,7 @@ def make_sync_call( ) completion_stream: Final = ModelResponseIterator( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, json_mode=json_mode, speed=speed, diff --git a/litellm/llms/anthropic/messages/caller_model_checks.py b/litellm/llms/anthropic/messages/caller_model_checks.py index 7eb6a8f9266..3822f075ccf 100644 --- a/litellm/llms/anthropic/messages/caller_model_checks.py +++ b/litellm/llms/anthropic/messages/caller_model_checks.py @@ -268,6 +268,7 @@ async def caller_within_model_budget( user_id=user_id, user_model_max_budget=user_model_max_budget, model=model, + team_id=user_api_key_auth.team_id, ) except litellm.BudgetExceededError: return "denied" @@ -309,6 +310,7 @@ async def caller_within_model_budget( end_user_id=end_user_id, end_user_model_max_budget=end_user_model_max_budget, model=model, + team_id=user_api_key_auth.team_id, ) except litellm.BudgetExceededError: return "denied" diff --git a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py index 614a7c96848..44f233f2254 100644 --- a/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/messages/responses_adapters/streaming_iterator.py @@ -2,15 +2,16 @@ ## Translates OpenAI call to Anthropic `/v1/messages` format import asyncio from collections import deque -from collections.abc import AsyncGenerator, Iterator, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final +import anyio from pydantic import ConfigDict, field_validator from litellm import verbose_logger from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid -from litellm.exceptions import MidStreamFallbackError +from litellm.exceptions import APIError, MidStreamFallbackError from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) @@ -22,6 +23,7 @@ from litellm.llms.anthropic.messages.native.utils import ( from litellm.llms.anthropic.messages.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.messages.utils import anthropic_sse_frame from litellm.responses.streaming_iterator import ( + map_stream_error_to_exception, partial_responses_incomplete_event_on_disconnect, stream_error_status_and_message, ) @@ -130,10 +132,12 @@ class AnthropicResponsesStreamWrapper: ) -> None: self.responses_stream = responses_stream self.model = model + self.litellm_logging_obj = litellm_logging_obj self.safeguards_evaluator: Final = safeguards_evaluator self._message_id: str = f"msg_{uuid.uuid4()}" if litellm_logging_obj is not None: litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id) + self._unreported_failure: Exception | None = None self._current_block_index: int = -1 # Map item_id -> content_block_index so we can stop the right block later self._item_id_to_block_index: dict[str, int] = {} @@ -377,16 +381,24 @@ class AnthropicResponsesStreamWrapper: if event_type == "response.failed": failed: Final = _FailedResponseEvent.model_validate(event) - status_code, message = stream_error_status_and_message( - failed.response.error if failed.response is not None else None - ) + upstream_error: Final = failed.response.error if failed.response is not None else None + status_code, message = stream_error_status_and_message(upstream_error) verbose_logger.error( "AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed (%s): %s", self.model, status_code, message, ) - self._fail_stream(status_code, message) + self._fail_stream( + status_code, + message, + failure=map_stream_error_to_exception( + upstream_error, + self.model, + getattr(self.responses_stream, "custom_llm_provider", None) or "", + logging_obj=self.litellm_logging_obj, + ), + ) return # ---- response completed -> message_delta + message_stop ---- @@ -439,10 +451,36 @@ class AnthropicResponsesStreamWrapper: self._sent_message_stop = True return - def _fail_stream(self, status_code: int, message: str) -> None: + def _fail_stream(self, status_code: int, message: str, failure: Exception | None = None) -> None: self._stream_failed = True + self._unreported_failure = failure or APIError( + status_code=status_code, + message=message, + llm_provider=getattr(self.responses_stream, "custom_llm_provider", None), + model=self.model, + ) self._chunk_queue.append(_anthropic_error_chunk(status_code, message)) + async def _report_unreported_failure(self) -> None: + failure: Final = self._unreported_failure + if failure is None: + return + self._unreported_failure = None + on_failure: Final[Callable[[Exception], Awaitable[None]] | None] = getattr( + self.litellm_logging_obj, "_on_detached_stream_failure", None + ) + if on_failure is None: + return + try: + with anyio.CancelScope(shield=True): + await on_failure(failure) + except Exception as hook_failure: # noqa: BLE001 # a failing proxy hook must not hide the error event from the client + verbose_logger.warning( + "AnthropicResponsesStreamWrapper: stream failure hook raised: %s(%s)", + type(hook_failure).__name__, + hook_failure, + ) + def __aiter__(self) -> "AnthropicResponsesStreamWrapper": return self @@ -450,6 +488,7 @@ class AnthropicResponsesStreamWrapper: if self._chunk_queue: return self._chunk_queue.popleft() if self._stream_failed: + await self._report_unreported_failure() raise StopAsyncIteration if not self._sent_message_start: @@ -478,7 +517,7 @@ class AnthropicResponsesStreamWrapper: verbose_logger.exception( "AnthropicResponsesStreamWrapper: upstream Responses stream for %s failed", self.model ) - self._fail_stream(*_failure_status_and_message(e)) + self._fail_stream(*_failure_status_and_message(e), failure=_original_failure(e)) if not self._chunk_queue and not self._sent_message_stop and not self._stream_failed: verbose_logger.error( @@ -498,8 +537,14 @@ class AnthropicResponsesStreamWrapper: async def _sse_frames(self) -> AsyncGenerator[bytes, None]: safeguard_results: Final = StreamedSafeguardResults(self.safeguards_evaluator) - async for chunk in self: - yield anthropic_sse_frame(await safeguard_results.observe(chunk)) + try: + async for chunk in self: + yield anthropic_sse_frame(await safeguard_results.observe(chunk)) + finally: + from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt + + if not closing_superseded_attempt(): + await self._report_unreported_failure() def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": return partial_responses_incomplete_event_on_disconnect(self.responses_stream) diff --git a/litellm/llms/anthropic/messages/stream_assembly.py b/litellm/llms/anthropic/messages/stream_assembly.py index 90bdad8f4bc..c103eb0f975 100644 --- a/litellm/llms/anthropic/messages/stream_assembly.py +++ b/litellm/llms/anthropic/messages/stream_assembly.py @@ -13,6 +13,7 @@ from litellm.llms.anthropic.chat.handler import ( ) from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.messages.native.utils import MESSAGE_DELTA_OPTIONAL_KEYS +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.utils import ( Choices, Message, @@ -269,7 +270,7 @@ def build_complete_streaming_response_legacy( def extract_sse_data(event_str: str) -> dict | None: """Parse the JSON object from the ``data:`` line of an Anthropic SSE event.""" - for line in event_str.splitlines(): + for line in split_sse_lines(event_str): stripped = line.strip() if stripped.startswith("data:"): payload = stripped[len("data:") :].strip() diff --git a/litellm/llms/anthropic/prompt_cache_prediction.py b/litellm/llms/anthropic/prompt_cache_prediction.py index 7a4c707860d..17c178d169e 100644 --- a/litellm/llms/anthropic/prompt_cache_prediction.py +++ b/litellm/llms/anthropic/prompt_cache_prediction.py @@ -29,7 +29,7 @@ from litellm.llms.anthropic.messages.native.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.anthropic.messages.native.utils import AnthropicMessagesRequestUtils, prepare_native_messages -from litellm.router_utils.baseline_request import ( +from litellm.router_strategy.auto_router.baseline_request import ( BASELINE_PARAMETERS, capture_baseline_parameters, ) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 02f1ba3472a..5a181faad11 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -13,13 +13,12 @@ from openai import ( BadRequestError, OpenAI, ) -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel import litellm -from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES +from litellm.constants import DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing -from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -34,7 +33,6 @@ from litellm.types.utils import ( from litellm.utils import ( CustomStreamWrapper, convert_to_model_response_object, - modify_url, ) from ...types.llms.openai import HttpxBinaryResponseContent @@ -59,8 +57,6 @@ from .image_generation import ( ) from .image_generation.http_utils import azure_deployment_image_generation_json_body -_DECODED_JSON: Final = TypeAdapter(object) - class AzureOpenAIAssistantsAPIConfig: """ @@ -860,11 +856,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers: dict, deployment_name: str | None = None, ) -> httpx.Response: - """ - Implemented for azure dall-e-2 image gen calls - - Alternative to needing a custom transport implementation - """ if client is None: _params: Final = {} if timeout is not None: @@ -881,95 +872,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): else: async_handler = client - if ( - "images/generations" in api_base - and api_version - in [ # dall-e-3 starts from `2023-12-01-preview` so we should be able to avoid conflict - "2023-06-01-preview", - "2023-07-01-preview", - "2023-08-01-preview", - "2023-09-01-preview", - "2023-10-01-preview", - ] - ): # CREATE + POLL for azure dall-e-2 calls - api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit") - - data.pop( - "model", None - ) # REMOVE 'model' from dall-e-2 arg https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-a-generated-image-dall-e-2-preview - response = await async_handler.post( - url=api_base, - data=json.dumps(data), - headers=headers, - ) - if "operation-location" in response.headers: - operation_location_url: Final = response.headers["operation-location"] - else: - raise AzureOpenAIError(status_code=500, message=response.text) - # Reject polling URLs that don't share an origin with ``api_base``. - # Without this an upstream-controlled or attacker-controlled - # value would receive the operator's Azure API key in the - # request headers below. VERIA-51. - try: - assert_same_origin(operation_location_url, api_base) - except SSRFError as ssrf_err: - raise AzureOpenAIError( - status_code=502, - message=f"Rejected polling URL: {ssrf_err}", - ) - response = await async_handler.get( - url=operation_location_url, - headers=headers, - ) - - await response.aread() - - timeout_secs: Final[int] = AZURE_OPERATION_POLLING_TIMEOUT - start_time: Final = time.time() - if "status" not in response.json(): - # Don't reflect the raw response body — when the polling - # URL points at an internal JSON API (cloud metadata - # service etc.) reflecting it here turns Blind SSRF into - # Full-Read SSRF. VERIA-51. - raise AzureOpenAIError( - status_code=502, - message="Polling response missing 'status' field", - ) - while response.json()["status"] not in ["succeeded", "failed"]: - if time.time() - start_time > timeout_secs: - raise AzureOpenAIError(status_code=408, message="Operation polling timed out.") - - await asyncio.sleep(int(response.headers.get("retry-after") or 10)) - response = await async_handler.get( - url=operation_location_url, - headers=headers, - ) - await response.aread() - - if response.json()["status"] == "failed": - error_data: Final = response.json() - # Preserve Azure error details (e.g. content_policy_violation, - # inner_error, content_filter_results) as structured body so - # exception_type() can route them correctly. - _error_body: Final = error_data.get("error", error_data) - _error_msg: Final = ( - _error_body.get("message", "Image generation failed") - if isinstance(_error_body, dict) - else json.dumps(error_data) - ) - raise AzureOpenAIError( - status_code=400, - message=_error_msg, - body=error_data, - ) - - result: Final = _DECODED_JSON.validate_python(response.json()["result"]) - return httpx.Response( - status_code=200, - headers=response.headers, - content=json.dumps(result).encode("utf-8"), - request=httpx.Request(method="POST", url="https://api.openai.com/v1"), - ) request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name) return await async_handler.post( url=api_base, @@ -988,11 +890,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers: dict, deployment_name: str | None = None, ) -> httpx.Response: - """ - Implemented for azure dall-e-2 image gen calls - - Alternative to needing a custom transport implementation - """ if client is None: _params: Final = {} if timeout is not None: @@ -1006,87 +903,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): else: sync_handler = client - if ( - "images/generations" in api_base - and api_version - in [ # dall-e-3 starts from `2023-12-01-preview` so we should be able to avoid conflict - "2023-06-01-preview", - "2023-07-01-preview", - "2023-08-01-preview", - "2023-09-01-preview", - "2023-10-01-preview", - ] - ): # CREATE + POLL for azure dall-e-2 calls - api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit") - - data.pop( - "model", None - ) # REMOVE 'model' from dall-e-2 arg https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-a-generated-image-dall-e-2-preview - response = sync_handler.post( - url=api_base, - data=json.dumps(data), - headers=headers, - ) - if "operation-location" in response.headers: - operation_location_url: Final = response.headers["operation-location"] - else: - raise AzureOpenAIError(status_code=500, message=response.text) - try: - assert_same_origin(operation_location_url, api_base) - except SSRFError as ssrf_err: - raise AzureOpenAIError( - status_code=502, - message=f"Rejected polling URL: {ssrf_err}", - ) - response = sync_handler.get( - url=operation_location_url, - headers=headers, - ) - - response.read() - - timeout_secs: Final[int] = AZURE_OPERATION_POLLING_TIMEOUT - start_time: Final = time.time() - if "status" not in response.json(): - raise AzureOpenAIError( - status_code=502, - message="Polling response missing 'status' field", - ) - while response.json()["status"] not in ["succeeded", "failed"]: - if time.time() - start_time > timeout_secs: - raise AzureOpenAIError(status_code=408, message="Operation polling timed out.") - - time.sleep(int(response.headers.get("retry-after") or 10)) - response = sync_handler.get( - url=operation_location_url, - headers=headers, - ) - response.read() - - if response.json()["status"] == "failed": - error_data: Final = response.json() - # Preserve Azure error details (e.g. content_policy_violation, - # inner_error, content_filter_results) as structured body so - # exception_type() can route them correctly. - _error_body: Final = error_data.get("error", error_data) - _error_msg: Final = ( - _error_body.get("message", "Image generation failed") - if isinstance(_error_body, dict) - else json.dumps(error_data) - ) - raise AzureOpenAIError( - status_code=400, - message=_error_msg, - body=error_data, - ) - - result: Final = _DECODED_JSON.validate_python(response.json()["result"]) - return httpx.Response( - status_code=200, - headers=response.headers, - content=json.dumps(result).encode("utf-8"), - request=httpx.Request(method="POST", url="https://api.openai.com/v1"), - ) request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name) return sync_handler.post( url=api_base, @@ -1193,10 +1009,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): deployment_name=model, ) - provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) + provider_config: Final = get_azure_image_generation_config(data.get("model", "")) if provider_config is not None: return provider_config.transform_image_generation_response( - model=data.get("model", "dall-e-2"), + model=data.get("model", ""), raw_response=httpx_response, model_response=model_response or ImageResponse(), logging_obj=logging_obj, @@ -1329,10 +1145,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers=request_headers, deployment_name=model, ) - provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) + provider_config: Final = get_azure_image_generation_config(data.get("model", "")) if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig): return provider_config.transform_image_generation_response( - model=data.get("model", "dall-e-2"), + model=data.get("model", ""), raw_response=httpx_response, model_response=model_response or ImageResponse(), logging_obj=logging_obj, diff --git a/litellm/llms/azure/image_generation/__init__.py b/litellm/llms/azure/image_generation/__init__.py index a2a905f2287..7e9b5f5d99c 100644 --- a/litellm/llms/azure/image_generation/__init__.py +++ b/litellm/llms/azure/image_generation/__init__.py @@ -4,14 +4,10 @@ from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from .dall_e_2_transformation import AzureDallE2ImageGenerationConfig -from .dall_e_3_transformation import AzureDallE3ImageGenerationConfig from .gpt_transformation import AzureGPTImageGenerationConfig from .http_utils import azure_deployment_image_generation_json_body __all__ = [ - "AzureDallE2ImageGenerationConfig", - "AzureDallE3ImageGenerationConfig", "AzureGPTImageGenerationConfig", "azure_deployment_image_generation_json_body", ] @@ -21,11 +17,7 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig: model = model.lower() model = model.replace("-", "") model = model.replace("_", "") - if model == "" or "dalle2" in model: # empty model is dall-e-2 - return AzureDallE2ImageGenerationConfig() - elif "dalle3" in model: - return AzureDallE3ImageGenerationConfig() - elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + if AzureFoundryMAIImageGenerationConfig.is_mai_model(model): return AzureFoundryMAIImageGenerationConfig() else: verbose_logger.debug( diff --git a/litellm/llms/azure/image_generation/dall_e_2_transformation.py b/litellm/llms/azure/image_generation/dall_e_2_transformation.py deleted file mode 100644 index 788753e4ea9..00000000000 --- a/litellm/llms/azure/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig - - -class AzureDallE2ImageGenerationConfig(DallE2ImageGenerationConfig): - """ - Azure dall-e-2 image generation config - """ diff --git a/litellm/llms/azure/image_generation/dall_e_3_transformation.py b/litellm/llms/azure/image_generation/dall_e_3_transformation.py deleted file mode 100644 index e6974f9317e..00000000000 --- a/litellm/llms/azure/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig - - -class AzureDallE3ImageGenerationConfig(DallE3ImageGenerationConfig): - """ - Azure dall-e-3 image generation config - """ diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 94f1399f01b..423a09c5c92 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -35,6 +35,7 @@ from litellm.llms.azure_ai.agents.transformation import ( AzureAIAgentsConfig, AzureAIAgentsError, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionAnnotationURLCitation, @@ -632,7 +633,7 @@ class AzureAIAgentsHandler: current_event = None - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): line = line.strip() if line.startswith("event:"): diff --git a/litellm/llms/azure_ai/image_generation/__init__.py b/litellm/llms/azure_ai/image_generation/__init__.py index 742922ad295..bd3c11d36f9 100644 --- a/litellm/llms/azure_ai/image_generation/__init__.py +++ b/litellm/llms/azure_ai/image_generation/__init__.py @@ -3,15 +3,11 @@ from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) -from .dall_e_2_transformation import AzureFoundryDallE2ImageGenerationConfig -from .dall_e_3_transformation import AzureFoundryDallE3ImageGenerationConfig from .flux_transformation import AzureFoundryFluxImageGenerationConfig from .gpt_transformation import AzureFoundryGPTImageGenerationConfig from .mai_transformation import AzureFoundryMAIImageGenerationConfig __all__ = [ - "AzureFoundryDallE2ImageGenerationConfig", - "AzureFoundryDallE3ImageGenerationConfig", "AzureFoundryFluxImageGenerationConfig", "AzureFoundryGPTImageGenerationConfig", "AzureFoundryMAIImageGenerationConfig", @@ -22,11 +18,7 @@ def get_azure_ai_image_generation_config(model: str) -> BaseImageGenerationConfi model = model.lower() model = model.replace("-", "") model = model.replace("_", "") - if model == "" or "dalle2" in model: # empty model is dall-e-2 - return AzureFoundryDallE2ImageGenerationConfig() - elif "dalle3" in model: - return AzureFoundryDallE3ImageGenerationConfig() - elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model): + if AzureFoundryMAIImageGenerationConfig.is_mai_model(model): return AzureFoundryMAIImageGenerationConfig() elif "flux" in model: return AzureFoundryFluxImageGenerationConfig() diff --git a/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py b/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py deleted file mode 100644 index 2c8a5426628..00000000000 --- a/litellm/llms/azure_ai/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig - - -class AzureFoundryDallE2ImageGenerationConfig(DallE2ImageGenerationConfig): - """ - Azure dall-e-2 image generation config - """ diff --git a/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py b/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py deleted file mode 100644 index 3ceb9b405da..00000000000 --- a/litellm/llms/azure_ai/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,7 +0,0 @@ -from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig - - -class AzureFoundryDallE3ImageGenerationConfig(DallE3ImageGenerationConfig): - """ - Azure dall-e-3 image generation config - """ diff --git a/litellm/llms/base_llm/agents/transformation.py b/litellm/llms/base_llm/agents/transformation.py index 9d139b289c4..5ea6ea950d9 100644 --- a/litellm/llms/base_llm/agents/transformation.py +++ b/litellm/llms/base_llm/agents/transformation.py @@ -15,6 +15,7 @@ from collections.abc import Mapping import httpx from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, @@ -50,10 +51,9 @@ class BaseAgentsAPIConfig(ABC): @abstractmethod def transform_create_request( self, - name: str, - litellm_params: Mapping[str, object], - ) -> dict[str, object]: - """Map name + litellm_params to the provider's create-agent body.""" + create_request: AgentCreateRequest, + ) -> Mapping[str, object]: + """Map the create request to the provider's create-agent body.""" @abstractmethod def transform_create_response( diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 032951be74a..3afc685b3dd 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -96,7 +96,7 @@ class BaseModelResponseIterator: released (and a backend like vLLM aborts generation) when the stream is abandoned before its natural end. - ``streaming_response`` is usually a bare ``aiter_lines()`` generator + ``streaming_response`` is usually a bare ``aiter_sse_lines()`` generator that holds no reference to the response, so the handler that owns the response attaches it here after construction.""" if self.http_response is not None: diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 4346e214dc8..491fb1919b6 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod -from collections.abc import AsyncGenerator, Iterator, Mapping -from typing import TYPE_CHECKING, Any, Union +from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Union import httpx from openai.types.file_deleted import FileDeleted @@ -38,6 +39,22 @@ else: UserAPIKeyAuth = Any +def _routes_nowhere(_model_name: str) -> bool: + return False + + +@dataclass(frozen=True, slots=True) +class FileListModelScope: + model_mapping_keys: frozenset[str] + routes_stored_model_name: Callable[[str], bool] = _routes_nowhere + + def holds(self, model_mapping_keys: Iterable[str], stored: bool) -> bool: + mapped_keys: Final = frozenset(model_mapping_keys) + return not self.model_mapping_keys.isdisjoint(mapped_keys) or ( + stored and any(map(self.routes_stored_model_name, mapped_keys)) + ) + + class BaseFileUploadStream(ABC): """Re-iterable request body that yields an upload's bytes lazily. @@ -274,6 +291,8 @@ class BaseFileEndpoints(ABC): user_api_key_dict: UserAPIKeyAuth, limit: int | None = None, after: str | None = None, + model_scope: FileListModelScope | None = None, + before: str | None = None, **data: dict, ) -> FileListPage: pass diff --git a/litellm/llms/base_llm/guardrail_translation/utils.py b/litellm/llms/base_llm/guardrail_translation/utils.py index 725146fe894..fb886e42780 100644 --- a/litellm/llms/base_llm/guardrail_translation/utils.py +++ b/litellm/llms/base_llm/guardrail_translation/utils.py @@ -6,6 +6,7 @@ from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row ha from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage @@ -24,7 +25,7 @@ def _anthropic_stream_chunk_events(item: object) -> list[dict]: events: Final[list[dict]] = [] for block in chunk.split("\n\n"): - for line in block.splitlines(): + for line in split_sse_lines(block): stripped = line.strip() if not stripped.startswith("data:"): continue diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py index fe0e3abb87e..9b9b4e0a947 100644 --- a/litellm/llms/base_llm/sandbox/transformation.py +++ b/litellm/llms/base_llm/sandbox/transformation.py @@ -12,6 +12,7 @@ from typing import Final import httpx from pydantic import Field, PrivateAttr +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.types.llms.base import LiteLLMPydanticObjectBase SANDBOX_MAX_OUTPUT_BYTES: Final = 10 * 1024 * 1024 @@ -99,7 +100,7 @@ class BaseSandboxConfig: async def _read_capped_lines(self, response: httpx.Response) -> list[str]: lines: Final[list[str]] = [] total = 0 - async for line in response.aiter_lines(): + async for line in aiter_sse_lines(response.aiter_text()): total += len(line.encode("utf-8")) if total > SANDBOX_MAX_OUTPUT_BYTES: raise ValueError( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index f7566ac4923..dc14ef005b5 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -10,6 +10,8 @@ import time from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -34,6 +36,15 @@ else: LiteLLMLoggingObj = Any +@with_config(ConfigDict(extra="allow", strict=True)) +class _PegasusResponse(TypedDict): + message: ReadOnly[NotRequired[str | None]] + finishReason: ReadOnly[NotRequired[str | None]] + + +_PEGASUS_RESPONSE: Final = TypeAdapter(_PegasusResponse) + + class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): """ Handles transforming OpenAI-style requests into Bedrock InvokeModel requests for @@ -207,7 +218,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): ModelResponse with choices[0].message.content and finish_reason """ try: - completion_response: Final = raw_response.json() + completion_response: Final = _PEGASUS_RESPONSE.validate_python(raw_response.json()) except Exception as e: raise BedrockError( message=f"Error parsing response: {raw_response.text}, error: {e}", @@ -224,7 +235,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): message_content: Final = completion_response.get("message", "") # Extract finish reason and map to LiteLLM format - finish_reason_raw: Final = completion_response.get("finishReason", "stop") + finish_reason_raw: Final = completion_response.get("finishReason") or "stop" finish_reason: Final = map_finish_reason(finish_reason_raw) # Set the response content diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 51b5a99272e..852179e8caa 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -1,11 +1,12 @@ import copy import json import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast, get_args import httpx -from pydantic import TypeAdapter, ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -48,6 +49,27 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM _GUARDRAIL_CONFIG_VALIDATOR: Final["TypeAdapter[GuardrailConfigBlock]"] = TypeAdapter(GuardrailConfigBlock) + +@with_config(ConfigDict(extra="allow", strict=True)) +class _CohereGeneration(TypedDict): + text: ReadOnly[str | None] + finish_reason: ReadOnly[str | None] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _CohereInvokeResponse(TypedDict): + text: NotRequired[ReadOnly[str | None]] + generations: NotRequired[ReadOnly[Sequence[_CohereGeneration]]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _LlamaInvokeResponse(TypedDict): + generation: ReadOnly[str | None] + + +_COHERE_INVOKE_RESPONSE: Final = TypeAdapter(_CohereInvokeResponse) +_LLAMA_INVOKE_RESPONSE: Final = TypeAdapter(_LlamaInvokeResponse) + _GUARDRAIL_CONFIG_EXPECTED_FORMAT: Final = ( "{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}" ) @@ -310,12 +332,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): outputText: str | None = None try: if provider == "cohere": - if "text" in completion_response: - outputText = completion_response["text"] - elif "generations" in completion_response: - outputText = completion_response["generations"][0]["text"] + cohere_response: Final = _COHERE_INVOKE_RESPONSE.validate_python(completion_response) + if "text" in cohere_response: + outputText = cohere_response["text"] + elif "generations" in cohere_response: + outputText = cohere_response["generations"][0]["text"] model_response.choices[0].finish_reason = map_finish_reason( - completion_response["generations"][0]["finish_reason"] + cohere_response["generations"][0]["finish_reason"] or "stop" ) elif provider == "anthropic": return litellm.AmazonAnthropicClaudeConfig().transform_response( @@ -361,7 +384,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): elif provider == "ai21": outputText = completion_response.get("completions")[0].get("data").get("text") elif provider == "meta" or provider == "llama" or provider == "deepseek_r1": - outputText = completion_response["generation"] + outputText = _LLAMA_INVOKE_RESPONSE.validate_python(completion_response)["generation"] elif provider == "mistral": outputText = litellm.AmazonMistralConfig.get_outputText(completion_response, model_response) else: # amazon titan diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 8c4702bcf99..0624c2c65db 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -49,6 +49,7 @@ from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsAuthParams, BedrockBatchRecordKind from litellm.types.llms.openai import ( @@ -400,7 +401,7 @@ def _batch_input_line(line_number: int, line: str) -> _BatchInputLine: def _non_blank_lines(file_content: str) -> Iterator[tuple[int, str]]: - numbered_lines: Final = enumerate(file_content.splitlines(), start=1) + numbered_lines: Final = enumerate(split_sse_lines(file_content), start=1) return ((line_number, line) for line_number, line in numbered_lines if line.strip()) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 1b01b0ec6d0..3a65a8928bf 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -45,6 +45,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.structured_edits imp structured_edits, tool_edits, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.guardrails.auto_router_compression import is_compression_guardrail from litellm.types.utils import GenericGuardrailAPIInputs @@ -947,7 +948,7 @@ def _invoke_frames(items: Sequence[object]) -> tuple[InvokeStreamFrame, ...]: def _sse_events(sse: bytes) -> tuple[Mapping[str, object], ...]: return tuple( _JSON_OBJECT.validate_python(json.loads(line[len("data:") :])) - for line in sse.decode().splitlines() + for line in split_sse_lines(sse.decode()) if line.startswith("data:") ) diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 9cb38822881..952d8858f78 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -53,6 +53,7 @@ from litellm.types.realtime import ( RealtimeResponseTypedDict, ) +_JSON_OBJECT_MESSAGE: Final = TypeAdapter(dict[str, JsonValue], config=ConfigDict(strict=True)) _JSON_VALUE: Final = TypeAdapter(object) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) @@ -856,7 +857,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final[dict[str, JsonValue]] = json.loads(message) + json_message: Final = _JSON_OBJECT_MESSAGE.validate_python(json.loads(message)) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] diff --git a/litellm/llms/bedrock/search/transformation.py b/litellm/llms/bedrock/search/transformation.py index 565fc12aa1e..6f72cadad67 100644 --- a/litellm/llms/bedrock/search/transformation.py +++ b/litellm/llms/bedrock/search/transformation.py @@ -47,6 +47,7 @@ from litellm.llms.base_llm.search.transformation import ( ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.secret_managers.main import get_secret_str # AgentCore web-search rejects queries longer than 200 characters @@ -146,7 +147,9 @@ def _iter_sse_events(text: str) -> Iterator[Mapping[str, object]]: progress notifications before the JSON-RPC response. """ for chunk in _SSE_EVENT_SEPARATOR.split(text): - payload = "\n".join(line[len("data:") :].lstrip() for line in chunk.splitlines() if line.startswith("data:")) + payload = "\n".join( + line[len("data:") :].lstrip() for line in split_sse_lines(chunk) if line.startswith("data:") + ) if not payload: continue try: diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 040b7b8c1d8..72b5cde21d1 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -11,6 +11,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo safe_convert_created_field, ) from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.sse_output_recovery import ( @@ -199,7 +200,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): error_message = None streamed_output_items: Final[dict[int, dict[str, object]]] = {} text_only_output_items: Final[dict[int, dict[str, object]]] = {} - for chunk in body_text.splitlines(): + for chunk in split_sse_lines(body_text): parsed_chunk = parse_sse_json_chunk(chunk) if parsed_chunk is None: continue diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 6bebfe2d444..6b172e77295 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.utils import TextChoices from litellm.utils import CustomStreamWrapper, TextCompletionResponse @@ -109,7 +110,7 @@ async def make_call( if response.status_code != 200: raise TextCompletionCodestralError(status_code=response.status_code, message=response.text) - completion_stream: Final = response.aiter_lines() + completion_stream: Final = aiter_sse_lines(response.aiter_text()) # LOGGING logging_obj.post_call( input=messages, @@ -339,13 +340,13 @@ class CodestralTextCompletion: response = litellm.module_level_client.post( completion_url, headers=headers, - data=json.dumps(data), + data=json.dumps({**data, "stream": True}), stream=stream, ) _response: Final = CustomStreamWrapper( - response.iter_lines(), + iter_sse_lines(response.iter_text()), model, - custom_llm_provider="codestral", + custom_llm_provider="text-completion-codestral", logging_obj=logging_obj, ) return _response diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 447d7b685f7..e6544968d24 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, get_ssl_configuration, ) +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.types.llms.openai import FileTypes from litellm.types.utils import HttpHandlerRequestFields, ImageResponse, LlmProviders from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager @@ -466,7 +467,7 @@ class BaseLLMAIOHTTPHandler: ) else: completion_stream = provider_config.get_model_response_iterator( - streaming_response=response.iter_lines(), sync_stream=True + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True ) # LOGGING @@ -560,11 +561,8 @@ class BaseLLMAIOHTTPHandler: organization: str | None = None, headers: dict | None = None, ) -> ImageResponse: - if model is None: - raise ValueError("model is required for non-openai image variations") - provider_config: Final = ProviderConfigManager.get_provider_image_variation_config( - model=model, # openai defaults to dall-e-2 + model=model or "", provider=LlmProviders(custom_llm_provider), ) diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 0e4d97fd770..9297e76624d 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -6,6 +6,7 @@ import ssl import sys import typing import urllib.request +import weakref from collections.abc import Callable, Generator from typing import ClassVar, Final @@ -109,6 +110,24 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]: raise mapped_exc(message) from exc +def _release_transport_anchor(_transport: object) -> None: + """Release a transport anchored to a response stream. See ``_anchor_transport_to``.""" + + +def _anchor_transport_to(stream: "AiohttpResponseStream", transport: object) -> None: + """Keep the transport alive for as long as the stream can still read. + + The transport finalizer closes the session it owns, and nothing on the path from a + live response back to that session references the transport, so a client dropped + mid-stream would close the session under the body. ``weakref.finalize`` holds the + transport in its own registry rather than on the stream, which keeps the transport + out of the response's reference cycle (an exception carrying the response, for one) + so it is finalized by refcount once the anchor drops and can still schedule an async + close, instead of being collected in the same cyclic pass as the session it closes. + """ + weakref.finalize(stream, _release_transport_anchor, transport) + + class AiohttpResponseStream(httpx.AsyncByteStream): CHUNK_SIZE = 1024 * 16 @@ -463,10 +482,12 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Re-raise if it's a different RuntimeError raise + stream: Final = AiohttpResponseStream(response) + _anchor_transport_to(stream, self) return httpx.Response( status_code=response.status, headers=response.headers, - stream=AiohttpResponseStream(response), + stream=stream, request=request, ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index afa12b58b29..e7da6586b97 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -118,6 +118,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, header_value, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, @@ -1086,7 +1087,7 @@ class BaseLLMHTTPHandler: completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator_for_request( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), @@ -1229,7 +1230,7 @@ class BaseLLMHTTPHandler: completion_stream: object = MockResponseIterator(model_response=model_response, json_mode=json_mode) else: completion_stream = provider_config.get_model_response_iterator_for_request( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, json_mode=json_mode, optional_params=_OPTIONAL_PARAMS.validate_python(optional_params), diff --git a/litellm/llms/custom_httpx/sse_lines.py b/litellm/llms/custom_httpx/sse_lines.py new file mode 100644 index 00000000000..71ea7d621eb --- /dev/null +++ b/litellm/llms/custom_httpx/sse_lines.py @@ -0,0 +1,74 @@ +""" +Split a decoded Server-Sent Events body into lines the way the SSE spec does. + +httpx's `iter_lines()` / `aiter_lines()` follow `str.splitlines()`, which also breaks on U+0085, U+2028, +U+2029, \\x0b, \\x0c and \\x1c-\\x1e. JSON leaves U+0085, U+2028 and U+2029 unescaped inside strings, so a +`data:` line carrying one of them gets cut in half. SSE lines end only at CRLF, LF or CR +(https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream) +""" + +import io +import re +from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator, Sequence +from typing import Final + +_SSE_LINE_BREAK: Final = re.compile(r"\r\n|\r|\n") +_SSE_LINE_WITH_BREAK: Final = re.compile(r"[^\r\n]*(?:\r\n|\r|\n)|[^\r\n]+") + + +def _split_lines(text: str) -> Sequence[str]: + if "\r" not in text: + return text.split("\n") + crlf_folded: Final = text.replace("\r\n", "\n") + if "\r" not in crlf_folded: + return crlf_folded.split("\n") + return _SSE_LINE_BREAK.split(text) + + +class _SSELineSplitter: + def __init__(self) -> None: + self._partial_line: io.StringIO = io.StringIO() + self._after_cr: bool = False + + def feed(self, text: str) -> tuple[str, ...]: + if not text: + return () + fresh: Final = text[1:] if self._after_cr and text.startswith("\n") else text + self._after_cr = fresh.endswith("\r") + *ended, tail = _split_lines(fresh) + if not ended: + _ = self._partial_line.write(tail) + return () + first: Final = self._partial_line.getvalue() + ended[0] + self._partial_line = io.StringIO() + _ = self._partial_line.write(tail) + return (first, *ended[1:]) + + def flush(self) -> tuple[str, ...]: + rest: Final = self._partial_line.getvalue() + self._partial_line = io.StringIO() + return (rest,) if rest else () + + +def iter_sse_lines(text_chunks: Iterable[str]) -> Iterator[str]: + splitter: Final = _SSELineSplitter() + for text in text_chunks: + yield from splitter.feed(text) + yield from splitter.flush() + + +def split_sse_lines(text: str) -> tuple[str, ...]: + return tuple(iter_sse_lines((text,))) + + +def split_sse_lines_keepends(text: str) -> tuple[str, ...]: + return tuple(_SSE_LINE_WITH_BREAK.findall(text)) + + +async def aiter_sse_lines(text_chunks: AsyncIterable[str]) -> AsyncIterator[str]: + splitter: Final = _SSELineSplitter() + async for text in text_chunks: + for line in splitter.feed(text): + yield line + for line in splitter.flush(): + yield line diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index d0c2b4dff0c..35132346261 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -14,6 +14,8 @@ Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import ensure_api_route @@ -28,6 +30,23 @@ from ..common_utils import DashScopeError DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1" +@with_config(ConfigDict(extra="allow", strict=True)) +class _DashScopeUsage(TypedDict, total=False): + prompt_tokens: NotRequired[ReadOnly[int]] + total_tokens: NotRequired[ReadOnly[int]] + + +@with_config(ConfigDict(extra="allow", strict=True)) +class _DashScopeEmbeddingResponse(TypedDict, total=False): + data: NotRequired[ReadOnly[list[object]]] + model: NotRequired[ReadOnly[str | None]] + usage: NotRequired[ReadOnly[_DashScopeUsage | None]] + id: NotRequired[ReadOnly[object]] + + +_DASHSCOPE_EMBEDDING_RESPONSE: Final = TypeAdapter(_DashScopeEmbeddingResponse) + + class DashScopeEmbeddingConfig(BaseEmbeddingConfig): """ Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api @@ -152,11 +171,12 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): message=message, ) + parsed_response: Final = _DASHSCOPE_EMBEDDING_RESPONSE.validate_python(response_json) model_response.object = "list" - model_response.data = response_json.get("data", []) - model_response.model = response_json.get("model", model) + model_response.data = parsed_response.get("data", []) + model_response.model = parsed_response.get("model", model) - usage: Final = response_json.get("usage") or {} + usage: Final = parsed_response.get("usage") or {} prompt_tokens: Final = usage.get("prompt_tokens", 0) total_tokens: Final = usage.get("total_tokens", prompt_tokens) setattr( @@ -169,8 +189,8 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): ), ) - if "id" in response_json: - setattr(model_response, "id", response_json["id"]) + if "id" in parsed_response: + setattr(model_response, "id", parsed_response["id"]) return model_response diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index eca87f9dd5f..cc85012b803 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -1,6 +1,8 @@ import json from typing import Final +from pydantic import ConfigDict, TypeAdapter + import litellm from litellm import verbose_logger from litellm.types.llms.openai import ( @@ -23,6 +25,9 @@ def _usage_block(usage: Usage) -> ChatCompletionUsageBlock: ) +_JSON_CHUNK: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + + class ModelResponseIterator: def __init__(self, streaming_response, sync_stream: bool): self.streaming_response = streaming_response @@ -99,7 +104,7 @@ class ModelResponseIterator: self.response_iterator = self.streaming_response chunk: Final = self._next_sse_data() try: - json_chunk: Final = json.loads(chunk) + json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk)) return self.chunk_parser(chunk=json_chunk) except ValueError as e: verbose_logger.debug( @@ -126,7 +131,7 @@ class ModelResponseIterator: if chunk == "[DONE]": raise StopAsyncIteration try: - json_chunk: Final = json.loads(chunk) + json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk)) return self.chunk_parser(chunk=json_chunk) except ValueError as e: verbose_logger.debug( diff --git a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py index 344d9d3c75b..6d0891295bd 100644 --- a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py +++ b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py @@ -4,13 +4,14 @@ from typing import Final from httpx._types import RequestFiles +from litellm.images.utils import image_to_data_url from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import FileTypes -from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig, to_data_url +from .transformation import DEFAULT_BASE_URL, FalAIImageEditConfig FLUX_LORA_DEPTH_ENDPOINT: Final[str] = "fal-ai/flux-lora-depth" SUPPORTED_OPENAI_PARAMS: Final[tuple[str, ...]] = ("n", "size") @@ -68,7 +69,7 @@ class FalAIFluxLoraDepthEditConfig(FalAIImageEditConfig): ) request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict "prompt": prompt, - "image_url": to_data_url(next(iter(images))), + "image_url": image_to_data_url(next(iter(images))), **provider_params, } return request_body, () diff --git a/litellm/llms/fal_ai/image_edit/transformation.py b/litellm/llms/fal_ai/image_edit/transformation.py index ea4e8cac5bf..5cbf635fd47 100644 --- a/litellm/llms/fal_ai/image_edit/transformation.py +++ b/litellm/llms/fal_ai/image_edit/transformation.py @@ -1,15 +1,12 @@ -import base64 -import os from collections.abc import Mapping -from pathlib import Path from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Final import httpx from httpx._types import RequestFiles from pydantic import ConfigDict, TypeAdapter -from litellm.images.utils import ImageEditRequestUtils +from litellm.images.utils import image_to_data_url from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.fal_ai.image_generation.gpt_image_2_transformation import ( map_gpt_image_quality, @@ -39,39 +36,6 @@ PARAM_TRANSLATION: Final[Mapping[str, str]] = MappingProxyType( ) -@runtime_checkable -class _SeekableBinaryReader(Protocol): - def tell(self) -> int: ... - - def seek(self, offset: int) -> int: ... - - def read(self) -> bytes: ... - - -def _read_image_bytes(image: object) -> bytes: - if isinstance(image, bytes): - return image - if isinstance(image, tuple): - return _read_image_bytes(image[1]) - if isinstance(image, os.PathLike): - return Path(image).read_bytes() - if isinstance(image, _SeekableBinaryReader): - position: Final = image.tell() - image.seek(0) - data: Final = image.read() - image.seek(position) - return data - raise ValueError(f"Unsupported image type for Fal AI image edit: {type(image).__name__}") - - -def to_data_url(image: object) -> str: - if isinstance(image, str): - return image - image_bytes: Final = _read_image_bytes(image) - mime_type: Final = ImageEditRequestUtils.get_image_content_type(image_bytes) - return f"data:{mime_type};base64,{base64.b64encode(image_bytes).decode('utf-8')}" - - def _first(value: object) -> object: return value[0] if isinstance(value, list) and value else value @@ -146,14 +110,14 @@ class FalAIImageEditConfig(BaseImageEditConfig): raise ValueError("Fal AI image edit requires at least one input image") mask: Final = _first(image_edit_optional_request_params.get("mask")) mask_field: Final[Mapping[str, str]] = ( - MappingProxyType({"mask_url": to_data_url(mask)}) if mask is not None else MappingProxyType({}) + MappingProxyType({"mask_url": image_to_data_url(mask)}) if mask is not None else MappingProxyType({}) ) provider_params: Final[Mapping[str, object]] = MappingProxyType( {key: value for key, value in image_edit_optional_request_params.items() if key != "mask"} ) request_body: Final[dict[str, object]] = { # mutable-ok: base class contract returns a dict "prompt": prompt, - "image_urls": tuple(to_data_url(img) for img in images), + "image_urls": tuple(image_to_data_url(img) for img in images), **mask_field, **provider_params, } diff --git a/litellm/llms/gemini/agents/transformation.py b/litellm/llms/gemini/agents/transformation.py index f7352da5fb6..45a6552a3fa 100644 --- a/litellm/llms/gemini/agents/transformation.py +++ b/litellm/llms/gemini/agents/transformation.py @@ -20,16 +20,13 @@ from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, AgentVersionsResponse, ) -# Keys inside litellm_params that should be forwarded to the Gemini -# create-agent body verbatim. -_GEMINI_AGENT_BODY_KEYS: Final = ("base_agent", "instructions", "base_environment") - class _JsonResponse(Protocol): def json(self) -> object: ... @@ -39,30 +36,6 @@ def _json_body(response: _JsonResponse) -> object: return response.json() -# LiteLLM-internal keys that must never be forwarded to Gemini. -_LITELLM_INTERNAL_KEYS: Final = frozenset( - { - "custom_llm_provider", - "api_key", - "api_base", - "make_public", - "cost_per_query", - "input_cost_per_token", - "output_cost_per_token", - "require_trace_id_on_calls_to_agent", - "require_trace_id_on_calls_by_agent", - "max_iterations", - "max_budget_per_session", - "guardrails", - "is_public", - "agent_name", - "agent_id", - "agent_card_params", - "provider_agent_response", - } -) - - class GeminiAgentsConfig(BaseAgentsAPIConfig): """ Configuration for the Google AI Studio (Gemini) native Agents API. @@ -143,14 +116,9 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): def transform_create_request( self, - name: str, - litellm_params: Mapping[str, object], - ) -> dict[str, object]: - body: Final[dict[str, object]] = {"name": name} - for key in _GEMINI_AGENT_BODY_KEYS: - value = litellm_params.get(key) - if value is not None: - body[key] = value + create_request: AgentCreateRequest, + ) -> Mapping[str, object]: + body: Final[Mapping[str, object]] = create_request.model_dump(mode="json", exclude_none=True) verbose_logger.debug("GeminiAgentsConfig create body: %s", body) return body diff --git a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py index 34007bc8fd4..b6b2aa6c28f 100644 --- a/litellm/llms/gemini/google_genai/guardrail_translation/handler.py +++ b/litellm/llms/gemini/google_genai/guardrail_translation/handler.py @@ -23,6 +23,7 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamTransformSink, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -121,7 +122,7 @@ def _parse_json_dict_or_none(payload: str) -> Mapping[str, object] | None: def _sse_payload_texts(sse_text: str) -> tuple[str, ...]: return tuple( text - for line in sse_text.splitlines() + for line in split_sse_lines(sse_text) if line.startswith("data:") for payload in (line[len("data:") :].strip(),) if payload and payload != "[DONE]" diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index 6f3e9a56b92..1895a97af7d 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -15,7 +15,7 @@ import httpx from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import process_response_headers +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, process_response_headers from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo @@ -41,6 +41,14 @@ else: _JsonObject: TypeAlias = dict[str, object] +def _has_agent_config_type(agent_config: object) -> bool: + match agent_config: + case {"type": str()}: + return True + case _: + return False + + class _InteractionPayload(TypedDict, total=False): """JSON body of an Interactions API interaction, keyed as ``InteractionsAPIResponse`` fields.""" @@ -105,6 +113,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): return [ "model", "agent", + "agent_config", "input", "tools", "system_instruction", @@ -188,6 +197,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): request_body["input"] = input optional_keys: Final = [ + "agent_config", "tools", "system_instruction", "stream", @@ -197,8 +207,14 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): "response_modalities", "previous_interaction_id", ] + dropped_keys: Final = ( + frozenset({"agent_config"}) + if drops_unsupported_params(litellm_params.get("drop_params")) + and (model or not _has_agent_config_type(optional_params.get("agent_config"))) + else frozenset() + ) for key in optional_keys: - if optional_params.get(key) is not None: + if key not in dropped_keys and optional_params.get(key) is not None: request_body[key] = optional_params[key] # response_mime_type is removed — fold it into response_format. diff --git a/litellm/llms/gigachat/authenticator.py b/litellm/llms/gigachat/authenticator.py index d43e203d7a1..f37c6e88896 100644 --- a/litellm/llms/gigachat/authenticator.py +++ b/litellm/llms/gigachat/authenticator.py @@ -12,6 +12,8 @@ from types import MappingProxyType from typing import Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.caching.caching import InMemoryCache @@ -37,6 +39,17 @@ _token_cache: Final = InMemoryCache() _NO_LITELLM_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) +@with_config(ConfigDict(extra="allow", strict=True)) +class _GigaChatTokenResponse(TypedDict): + tok: NotRequired[ReadOnly[str | None]] + access_token: NotRequired[ReadOnly[str | None]] + exp: NotRequired[ReadOnly[int | float | str | None]] + expires_at: NotRequired[ReadOnly[int | float | str | None]] + + +_TOKEN_RESPONSE: Final = TypeAdapter(_GigaChatTokenResponse) + + class GigaChatAuthError(BaseLLMException): """GigaChat authentication error.""" @@ -230,7 +243,7 @@ async def _request_token_async( def _parse_token_response(response: httpx.Response) -> tuple[str, int]: """Parse OAuth token response.""" - data: Final = response.json() + data: Final = _TOKEN_RESPONSE.validate_python(response.json()) # GigaChat returns either 'tok'/'exp' or 'access_token'/'expires_at' access_token: Final = data.get("tok") or data.get("access_token") diff --git a/litellm/llms/huggingface/embedding/transformation.py b/litellm/llms/huggingface/embedding/transformation.py index 1b7d6483739..6e6a117227e 100644 --- a/litellm/llms/huggingface/embedding/transformation.py +++ b/litellm/llms/huggingface/embedding/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( ) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -357,7 +358,7 @@ class HuggingFaceEmbeddingConfig(BaseConfig): api_key: str | None = None, ) -> list[dict[str, str]]: streamed_response: Final = CustomStreamWrapper( - completion_stream=response.iter_lines(), + completion_stream=iter_sse_lines(response.iter_text()), model=model, custom_llm_provider="huggingface", logging_obj=logging_obj, diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 575820bceee..6fe355ad98c 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -12,6 +12,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices _DECODED_JSON: Final = TypeAdapter(object) @@ -36,12 +37,12 @@ class LangGraphSSEStreamIterator: def __iter__(self): """Initialize sync iteration.""" - self.line_iterator = self.response.iter_lines() + self.line_iterator = iter_sse_lines(self.response.iter_text()) return self def __aiter__(self): """Initialize async iteration.""" - self.async_line_iterator = self.response.aiter_lines() + self.async_line_iterator = aiter_sse_lines(self.response.aiter_text()) return self def _parse_sse_line(self, line: str) -> ModelResponseStream | None: diff --git a/litellm/llms/ollama/decisions/transformation.py b/litellm/llms/ollama/decisions/transformation.py new file mode 100644 index 00000000000..3e69f5cfede --- /dev/null +++ b/litellm/llms/ollama/decisions/transformation.py @@ -0,0 +1,17 @@ +from typing import Final + +from litellm.constants import PLACEHOLDER_API_KEY +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig + + +class OllamaDecisionsConfig(BaseDecisionsConfig): + api_key_env = ("OLLAMA_API_KEY",) + api_base_env = ("OLLAMA_API_BASE",) + api_key_required = False + + def get_default_api_base(self) -> str | None: + return "http://localhost:11434" + + def resolve_api_key(self, api_key: str | None) -> str | None: + resolved: Final = super().resolve_api_key(api_key) + return None if resolved == PLACEHOLDER_API_KEY else resolved diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index 8879c20c5eb..b4dc57ff5ed 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final import litellm from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import get_httpx_client +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines from litellm.utils import EmbeddingResponse, ModelResponse, Usage from ..common_utils import OobaboogaError @@ -74,7 +75,7 @@ def completion( stream=optional_params["stream"] if "stream" in optional_params else False, ) if "stream" in optional_params and optional_params["stream"] is True: - return response.iter_lines() + return iter_sse_lines(response.iter_text()) else: return oobabooga_config.transform_response( model=model, diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 97e988ef2b1..b7396e8fb03 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -76,6 +76,7 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _OBJECT_MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) +_CHOICE_MAPPINGS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) _DECODED_JSON: Final = TypeAdapter(object) @@ -958,8 +959,13 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): status_code=error_status_code, message=error_message, ) - choices = chunk.get("choices", []) - choices = self._map_reasoning_to_reasoning_content(choices) + raw_choices: Final = chunk.get("choices") or [] + if any(choice.get("finish_reason") == "abort" for choice in _CHOICE_MAPPINGS.validate_python(raw_choices)): + raise OpenAIError( + status_code=500, + message="Upstream generation aborted (finish_reason=abort)", + ) + choices: Final = self._map_reasoning_to_reasoning_content(raw_choices) fingerprint: Final = _OBJECT_MAPPING.validate_python(chunk).get("system_fingerprint") kwargs: Final[dict[str, Any]] = { diff --git a/litellm/llms/openai/image_edit/__init__.py b/litellm/llms/openai/image_edit/__init__.py index 63a41b65e8a..261010071c9 100644 --- a/litellm/llms/openai/image_edit/__init__.py +++ b/litellm/llms/openai/image_edit/__init__.py @@ -1,31 +1,5 @@ -from typing import Final - -from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig - -from .dalle2_transformation import DallE2ImageEditConfig from .transformation import OpenAIImageEditConfig __all__ = [ - "DallE2ImageEditConfig", "OpenAIImageEditConfig", - "get_openai_image_edit_config", ] - - -def get_openai_image_edit_config(model: str) -> BaseImageEditConfig: - """ - Get the appropriate OpenAI image edit config based on the model. - - Args: - model: The model name (e.g., "dall-e-2", "gpt-image-1") - - Returns: - The appropriate config instance for the model - """ - model_normalized: Final = model.lower().replace("-", "").replace("_", "") - - if model_normalized == "dalle2": - return DallE2ImageEditConfig() - else: - # Default to standard OpenAI config for gpt-image-1 and other models - return OpenAIImageEditConfig() diff --git a/litellm/llms/openai/image_edit/dalle2_transformation.py b/litellm/llms/openai/image_edit/dalle2_transformation.py deleted file mode 100644 index f349959b14f..00000000000 --- a/litellm/llms/openai/image_edit/dalle2_transformation.py +++ /dev/null @@ -1,98 +0,0 @@ -from io import BufferedReader -from typing import TYPE_CHECKING, Any, Final, cast - -from httpx._types import RequestFiles - -import litellm -from litellm.images.utils import ImageEditRequestUtils -from litellm.types.images.main import ImageEditRequestParams -from litellm.types.llms.openai import FileTypes -from litellm.types.router import GenericLiteLLMParams - -from .transformation import OpenAIImageEditConfig - -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj - - LiteLLMLoggingObj = _LiteLLMLoggingObj -else: - LiteLLMLoggingObj = Any - - -class DallE2ImageEditConfig(OpenAIImageEditConfig): - """ - DALL-E-2 specific configuration for image edit API. - - DALL-E-2 only supports editing a single image (not an array). - Uses "image" field name instead of "image[]". - """ - - def transform_image_edit_request( - self, - model: str, - prompt: str | None, - image: FileTypes | None, - image_edit_optional_request_params: dict, - litellm_params: GenericLiteLLMParams, - headers: dict, - ) -> tuple[dict, RequestFiles]: - """ - Transform image edit request for DALL-E-2. - - DALL-E-2 only accepts a single image with field name "image" (not "image[]"). - """ - request_params: Final = { - "model": model, - **image_edit_optional_request_params, - } - if image is not None: - request_params["image"] = image - if prompt is not None: - request_params["prompt"] = prompt - - request: Final = ImageEditRequestParams(**request_params) - request_dict: Final = cast(dict, request) - - ######################################################### - # Separate images and masks as `files` and send other parameters as `data` - ######################################################### - _image_list: Final = request_dict.get("image") - _mask = request_dict.get("mask") - data_without_files: Final = {k: v for k, v in request_dict.items() if k not in ["image", "mask"]} - files_list: Final[list[tuple[str, Any]]] = [] - - # Handle image parameter - DALL-E-2 only supports single image - if _image_list is not None: - image_list: Final = [_image_list] if not isinstance(_image_list, list) else _image_list - - # Validate only one image is provided - if len(image_list) > 1: - raise litellm.BadRequestError( - message="DALL-E-2 only supports editing a single image. Please provide one image.", - model=model, - llm_provider="openai", - ) - - # Use "image" field name (singular) for DALL-E-2 - for _image in image_list: - if _image is not None: - self._add_image_to_files( - files_list=files_list, - image=_image, - field_name="image", - ) - - # Handle mask parameter if provided - if _mask is not None: - # Handle case where mask can be a list (extract first mask) - if isinstance(_mask, list): - _mask = _mask[0] if _mask else None - - if _mask is not None: - mask_content_type: Final[str] = ImageEditRequestUtils.get_image_content_type(_mask) - if isinstance(_mask, BufferedReader): - files_list.append(("mask", (_mask.name, _mask, mask_content_type))) - else: - files_list.append(("mask", ("mask.png", _mask, mask_content_type))) - - return data_without_files, files_list diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index 9def833052b..0bba58096ef 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -96,7 +96,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig): """ # Build request params, only including non-None values request_params: Final = { - "model": model, + **({"model": model} if model else {}), **image_edit_optional_request_params, } if image is not None: diff --git a/litellm/llms/openai/image_generation/__init__.py b/litellm/llms/openai/image_generation/__init__.py index e20c80f20bb..4cffbe81a0e 100644 --- a/litellm/llms/openai/image_generation/__init__.py +++ b/litellm/llms/openai/image_generation/__init__.py @@ -1,9 +1,3 @@ -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) - -from .dall_e_2_transformation import DallE2ImageGenerationConfig -from .dall_e_3_transformation import DallE3ImageGenerationConfig from .gpt_transformation import GPTImageGenerationConfig from .guardrail_translation import ( OpenAIImageGenerationHandler, @@ -11,18 +5,7 @@ from .guardrail_translation import ( ) __all__ = [ - "DallE2ImageGenerationConfig", - "DallE3ImageGenerationConfig", "GPTImageGenerationConfig", "OpenAIImageGenerationHandler", "guardrail_translation_mappings", ] - - -def get_openai_image_generation_config(model: str) -> BaseImageGenerationConfig: - if model.startswith("dall-e-2") or model == "": # empty model is dall-e-2 - return DallE2ImageGenerationConfig() - elif model.startswith("dall-e-3"): - return DallE3ImageGenerationConfig() - else: - return GPTImageGenerationConfig() diff --git a/litellm/llms/openai/image_generation/dall_e_2_transformation.py b/litellm/llms/openai/image_generation/dall_e_2_transformation.py deleted file mode 100644 index ffc6f1d5fe9..00000000000 --- a/litellm/llms/openai/image_generation/dall_e_2_transformation.py +++ /dev/null @@ -1,81 +0,0 @@ -from typing import TYPE_CHECKING, Final - -import httpx - -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams -from litellm.types.utils import ImageResponse -from litellm.utils import convert_to_model_response_object - -if TYPE_CHECKING: - from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj - - from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer - - -class DallE2ImageGenerationConfig(BaseImageGenerationConfig): - """ - OpenAI dall-e-2 image generation config - """ - - def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - return ["n", "response_format", "quality", "size", "user"] - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - supported_params: Final = self.get_supported_openai_params(model) - for k in non_default_params: - if k not in optional_params: - if k in supported_params: - optional_params[k] = non_default_params[k] - elif drop_params: - pass - else: - raise ValueError( - f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." - ) - - return optional_params - - def transform_image_generation_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ImageResponse, - logging_obj: "LiteLLMLoggingObj", - request_data: dict, - optional_params: dict, - litellm_params: dict, - encoding: "Tokenizer | None", - api_key: str | None = None, - json_mode: bool | None = None, - ) -> ImageResponse: - response: Final = raw_response.json() - - stringified_response: Final = response - ## LOGGING - logging_obj.post_call( - input=request_data.get("prompt", ""), - api_key=api_key, - additional_args={"complete_input_dict": request_data}, - original_response=stringified_response, - ) - image_response: Final[ImageResponse] = convert_to_model_response_object( - response_object=stringified_response, - model_response_object=model_response, - response_type="image_generation", - ) - - # set optional params - image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 - image_response.quality = optional_params.get("quality", "standard") # always standard for dall-e-2 - image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-2 - - return image_response diff --git a/litellm/llms/openai/image_generation/dall_e_3_transformation.py b/litellm/llms/openai/image_generation/dall_e_3_transformation.py deleted file mode 100644 index 90b7eaedf2f..00000000000 --- a/litellm/llms/openai/image_generation/dall_e_3_transformation.py +++ /dev/null @@ -1,81 +0,0 @@ -from typing import TYPE_CHECKING, Final - -import httpx - -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams -from litellm.types.utils import ImageResponse -from litellm.utils import convert_to_model_response_object - -if TYPE_CHECKING: - from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj - - from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer - - -class DallE3ImageGenerationConfig(BaseImageGenerationConfig): - """ - OpenAI dall-e-3 image generation config - """ - - def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]: - return ["n", "response_format", "quality", "size", "user", "style"] - - def map_openai_params( - self, - non_default_params: dict, - optional_params: dict, - model: str, - drop_params: bool, - ) -> dict: - supported_params: Final = self.get_supported_openai_params(model) - for k in non_default_params: - if k not in optional_params: - if k in supported_params: - optional_params[k] = non_default_params[k] - elif drop_params: - pass - else: - raise ValueError( - f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters." - ) - - return optional_params - - def transform_image_generation_response( - self, - model: str, - raw_response: httpx.Response, - model_response: ImageResponse, - logging_obj: "LiteLLMLoggingObj", - request_data: dict, - optional_params: dict, - litellm_params: dict, - encoding: "Tokenizer | None", - api_key: str | None = None, - json_mode: bool | None = None, - ) -> ImageResponse: - response: Final = raw_response.json() - - stringified_response: Final = response - ## LOGGING - logging_obj.post_call( - input=request_data.get("prompt", ""), - api_key=api_key, - additional_args={"complete_input_dict": request_data}, - original_response=stringified_response, - ) - image_response: Final[ImageResponse] = convert_to_model_response_object( - response_object=stringified_response, - model_response_object=model_response, - response_type="image_generation", - ) - - # set optional params - image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024 - image_response.quality = optional_params.get("quality", "hd") # always hd for dall-e-3 - image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-3 - - return image_response diff --git a/litellm/llms/openai/image_variations/handler.py b/litellm/llms/openai/image_variations/handler.py index 89dcdc40e02..01f01d3de1d 100644 --- a/litellm/llms/openai/image_variations/handler.py +++ b/litellm/llms/openai/image_variations/handler.py @@ -151,7 +151,7 @@ class OpenAIImageVariationsHandler: ) -> ImageResponse: try: provider_config: Final = ProviderConfigManager.get_provider_image_variation_config( - model=model or "", # openai defaults to dall-e-2 + model=model or "", provider=LlmProviders.OPENAI, ) diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index 080fc3d24c9..b7e4cc6868f 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -5,7 +5,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. """ import ssl -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Final, cast from litellm._logging import redact_string, verbose_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES @@ -23,6 +23,9 @@ from ....llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..common_utils import is_openai_backed_api_base from ..openai import OpenAIChatCompletion +if TYPE_CHECKING: + from fastapi import WebSocket + class OpenAIRealtime(OpenAIChatCompletion): """ @@ -131,7 +134,7 @@ class OpenAIRealtime(OpenAIChatCompletion): async def async_realtime( self, model: str, - websocket: Any, + websocket: "WebSocket", logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 2c74ef19e7d..a450947afdf 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -16,6 +16,7 @@ from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.llms.databricks.streaming_utils import ModelResponseIterator from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openai.openai import OpenAIConfig @@ -65,7 +66,9 @@ async def make_call( model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: - completion_stream = ModelResponseIterator(streaming_response=response.aiter_lines(), sync_stream=False) + completion_stream = ModelResponseIterator( + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False + ) # LOGGING logging_obj.post_call( input=messages, @@ -103,7 +106,9 @@ def make_sync_call( model_response: Final = _fake_streamed_model_response(response.json()) completion_stream = MockResponseIterator(model_response=model_response) else: - completion_stream = ModelResponseIterator(streaming_response=response.iter_lines(), sync_stream=True) + completion_stream = ModelResponseIterator( + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True + ) # LOGGING logging_obj.post_call( diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index 0034ed2b9d4..2b71707239b 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -58,6 +58,12 @@ "max_completion_tokens": "max_tokens" } }, + "zerogpu": { + "base_url": "https://api.zerogpu.ai/v1", + "api_key_env": "ZEROGPU_API_KEY", + "api_base_env": "ZEROGPU_API_BASE", + "supported_endpoints": ["/v1/chat/completions", "/v1/responses"] + }, "chutes": { "base_url": "https://llm.chutes.ai/v1/", "api_key_env": "CHUTES_API_KEY", diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py index 7b4da933976..147d344795d 100644 --- a/litellm/llms/perplexity/embedding/transformation.py +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -12,12 +12,14 @@ This module decodes them into float arrays for OpenAI-compatible responses. """ import base64 +import builtins import struct -from collections.abc import Mapping -from typing import Final +from collections.abc import Mapping, Sequence +from typing import Final, Literal import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.url_utils import join_api_base @@ -31,6 +33,17 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i _INT: Final = TypeAdapter(int) +@with_config(ConfigDict(extra="allow", strict=True)) +class _PerplexityEmbeddingResponse(TypedDict, total=False): + model: ReadOnly[str | None] + object: ReadOnly[Literal["list"]] + data: ReadOnly[Sequence[builtins.object] | Mapping[str, builtins.object] | Literal[""]] + usage: ReadOnly[builtins.object] + + +_EMBEDDING_RESPONSE_ADAPTER: Final = TypeAdapter(_PerplexityEmbeddingResponse) + + class PerplexityEmbeddingError(BaseLLMException): def __init__( self, @@ -153,14 +166,18 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict = {}, ) -> EmbeddingResponse: try: - raw_response_json: Final = raw_response.json() + raw_response_json: Final[object] = raw_response.json() except Exception: raise PerplexityEmbeddingError(message=raw_response.text, status_code=raw_response.status_code) - model_response.model = raw_response_json.get("model", model) - model_response.object = raw_response_json.get("object", "list") + response_data: Final[_PerplexityEmbeddingResponse] = _EMBEDDING_RESPONSE_ADAPTER.validate_python( + raw_response_json + ) - raw_data: Final = raw_response_json.get("data", []) + model_response.model = response_data.get("model", model) + model_response.object = response_data.get("object", "list") + + raw_data: Final = response_data.get("data", []) decoded_data: Final[list[dict[str, object]]] = [] for item in raw_data: decoded_item = dict(_JSON_OBJECT.validate_python(item)) @@ -168,7 +185,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): decoded_data.append(decoded_item) model_response.data = decoded_data - usage_data: Final = _JSON_OBJECT.validate_python(raw_response_json.get("usage", {})) + usage_data: Final = _JSON_OBJECT.validate_python(response_data.get("usage", {})) usage: Final = Usage( prompt_tokens=_INT.validate_python( usage_data.get("prompt_tokens", 0) or usage_data.get("total_tokens", 0) or 0 diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py index d0a61ea6e00..2ba4a63e08b 100644 --- a/litellm/llms/predibase/chat/handler.py +++ b/litellm/llms/predibase/chat/handler.py @@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.utils import CustomStreamWrapper, ModelResponse from ..common_utils import PredibaseError @@ -34,7 +35,7 @@ async def make_call( if response.status_code != 200: raise PredibaseError(status_code=response.status_code, message=response.text) - completion_stream: Final = response.aiter_lines() + completion_stream: Final = aiter_sse_lines(response.aiter_text()) # LOGGING logging_obj.post_call( input=messages, @@ -163,7 +164,7 @@ class PredibaseChatCompletion: timeout=timeout, ) _response: Final = CustomStreamWrapper( - response.iter_lines(), + iter_sse_lines(response.iter_text()), model, custom_llm_provider="predibase", logging_obj=logging_obj, diff --git a/litellm/llms/tencent/cost_calculator.py b/litellm/llms/tencent/cost_calculator.py index d9aebdc3284..5d2abe3acf4 100644 --- a/litellm/llms/tencent/cost_calculator.py +++ b/litellm/llms/tencent/cost_calculator.py @@ -1,6 +1,49 @@ -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.types.utils import Usage +from typing import Final + +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_image_response_cost_from_usage, + flat_image_cost, + generic_cost_per_token, + resolve_image_model_info, +) +from litellm.types.utils import ImageResponse, ModelInfo, Usage def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: return generic_cost_per_token(model=model, usage=usage, custom_llm_provider="tencent") + + +def _deployment_image_token_pricing(model_info: ModelInfo | None) -> ModelInfo | None: + if model_info is None or model_info.get("output_cost_per_image_token") is not None: + return model_info + output_cost_per_token: Final = model_info.get("output_cost_per_token") + if output_cost_per_token is None: + return model_info + priced: Final[ModelInfo] = {**model_info, "output_cost_per_image_token": output_cost_per_token} + return priced + + +def _deployment_prices_per_image_only(model_info: ModelInfo | None) -> bool: + return ( + model_info is not None + and model_info.get("output_cost_per_image") is not None + and model_info.get("output_cost_per_image_token") is None + and model_info.get("output_cost_per_token") is None + ) + + +def image_cost_calculator(model: str, image_response: ImageResponse, model_info: ModelInfo | None = None) -> float: + resolved_model_info: Final = resolve_image_model_info( + model=model, custom_llm_provider="tencent", model_info=_deployment_image_token_pricing(model_info) + ) + if _deployment_prices_per_image_only(model_info): + return flat_image_cost(resolved_model_info, image_response) + token_based_cost: Final = calculate_image_response_cost_from_usage( + model=model, + image_response=image_response, + custom_llm_provider="tencent", + model_info=resolved_model_info, + ) + if token_based_cost is not None: + return token_based_cost + return flat_image_cost(resolved_model_info, image_response) diff --git a/litellm/llms/tencent/image_edit/__init__.py b/litellm/llms/tencent/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/tencent/image_edit/transformation.py b/litellm/llms/tencent/image_edit/transformation.py new file mode 100644 index 00000000000..3033d161ce7 --- /dev/null +++ b/litellm/llms/tencent/image_edit/transformation.py @@ -0,0 +1,147 @@ +from collections.abc import Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeAlias + +import httpx +from httpx._types import RequestFiles +from pydantic import TypeAdapter +from typing_extensions import override + +from litellm.images.utils import image_to_data_url +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig +from litellm.llms.tencent.image_generation.transformation import ( + TOKENHUB_IMAGE_PARAMS, + tokenhub_image_headers, + tokenhub_image_params, + tokenhub_image_response, + tokenhub_image_url, +) +from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ImageResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +_IMAGES: Final = TypeAdapter(tuple[object, ...]) +_FORM_BOOLEANS: Final = MappingProxyType({"true": True, "false": False}) +_OPENAI_PARAMS: Final = frozenset(TOKENHUB_IMAGE_PARAMS) + +_FormEntries: TypeAlias = tuple[tuple[tuple[str, ...], object], ...] + + +def _form_scalar(value: object) -> object: + match value: + case str() if value.lower() in _FORM_BOOLEANS: + return _FORM_BOOLEANS[value.lower()] + case str() if value.removeprefix("-").isdigit(): + return int(value) + case _: + return value + + +def _form_path(key: str) -> tuple[str, ...]: + root, bracket, rest = key.partition("[") + if not bracket or not rest.endswith("]"): + return (key,) + return (root, *rest.removesuffix("]").split("][")) + + +def _nested_form_value(entries: _FormEntries) -> object: + match entries: + case (((), value),): + return _form_scalar(value) + case _: + return _nested_form_fields(tuple(entry for entry in entries if entry[0])) + + +def _nested_form_fields(entries: _FormEntries) -> Mapping[str, object]: + roots: Final = tuple(dict.fromkeys(path[0] for path, _ in entries)) + return { + root: _nested_form_value(tuple((path[1:], value) for path, value in entries if path[0] == root)) + for root in roots + } + + +def _tokenhub_form_params(params: Mapping[str, object]) -> Mapping[str, object]: + tokenhub_entries: Final = tuple( + (_form_path(key), value) for key, value in params.items() if key not in _OPENAI_PARAMS + ) + return { + **{key: value for key, value in params.items() if key in _OPENAI_PARAMS}, + **_nested_form_fields(tokenhub_entries), + } + + +class TencentImageEditConfig(BaseImageEditConfig): + @override + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: mirrors override contract + return list(TOKENHUB_IMAGE_PARAMS) + + @override + def map_openai_params( + self, + image_edit_optional_params: ImageEditOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: mirrors override contract + return dict(tokenhub_image_params(image_edit_optional_params, model, drop_params)) + + @override + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None = None, + litellm_params: Mapping[str, object] | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: mirrors override contract + return dict(tokenhub_image_headers(headers, api_key, model)) + + @override + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: Mapping[str, object], + ) -> str: + return tokenhub_image_url(api_base) + + @override + def use_multipart_form_data(self) -> bool: + return False + + @override + def transform_image_edit_request( + self, + model: str, + prompt: str | None, + image: object, + image_edit_optional_request_params: Mapping[str, object], + litellm_params: GenericLiteLLMParams, + headers: Mapping[str, str], + ) -> tuple[dict[str, object], RequestFiles]: # mutable-ok: mirrors override contract + images: Final = _IMAGES.validate_python(image) if isinstance(image, list) else (image,) + return { + "model": model, + "prompt": prompt, + "images": tuple(image_to_data_url(img) for img in images if img is not None), + **_tokenhub_form_params(image_edit_optional_request_params), + }, () + + @override + def transform_image_edit_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> ImageResponse: + image_response: Final = tokenhub_image_response(raw_response) + if image_response is None: + raise BaseLLMException( + status_code=raw_response.status_code, + message=raw_response.text, + headers=raw_response.headers, + ) + return image_response diff --git a/litellm/llms/tencent/image_generation/__init__.py b/litellm/llms/tencent/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/tencent/image_generation/transformation.py b/litellm/llms/tencent/image_generation/transformation.py new file mode 100644 index 00000000000..fb4b102e8be --- /dev/null +++ b/litellm/llms/tencent/image_generation/transformation.py @@ -0,0 +1,200 @@ +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Final + +import httpx +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from typing_extensions import override + +from litellm.exceptions import AuthenticationError, UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIImageGenerationOptionalParams, +) +from litellm.types.utils import ( + ImageObject, + ImageResponse, + ImageUsage, + ImageUsageInputTokensDetails, + LlmProviders, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +TOKENHUB_API_BASE: Final = "https://tokenhub-intl.tencentcloudmaas.com/v1" +TOKENHUB_IMAGE_PATH: Final = "/wand/si-image/generation" +TOKENHUB_IMAGE_MODEL_PREFIX: Final = "seedream-" +TOKENHUB_IMAGE_PARAMS: Final[tuple[OpenAIImageGenerationOptionalParams, ...]] = ( + "n", + "size", + "response_format", + "output_format", + "background", + "user", +) + +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) + + +class _TokenHubImage(BaseModel): + model_config = ConfigDict(frozen=True) + + url: str | None = None + b64_json: str | None = None + size: str | None = None + output_format: str | None = None + + +class _TokenHubUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + total_tokens: int + + +class _TokenHubImageResponse(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + + created: int + data: tuple[_TokenHubImage, ...] + tokenhub_usage: _TokenHubUsage | None = None + + +def is_tokenhub_image_model(model: str) -> bool: + return model.removeprefix(f"{LlmProviders.TENCENT.value}/").startswith(TOKENHUB_IMAGE_MODEL_PREFIX) + + +def tokenhub_image_url(api_base: str | None) -> str: + base: Final = api_base or get_secret_str("TENCENT_API_BASE") or TOKENHUB_API_BASE + return join_api_base(ensure_api_route(base, "v1"), TOKENHUB_IMAGE_PATH) + + +def tokenhub_image_headers(headers: Mapping[str, str], api_key: str | None, model: str) -> Mapping[str, str]: + resolved_api_key: Final = api_key or get_secret_str("TENCENT_API_KEY") + if not resolved_api_key: + raise AuthenticationError( + message="TENCENT_API_KEY is not set. Set it in the environment or pass api_key.", + llm_provider="tencent", + model=model, + ) + return {**headers, "Authorization": f"Bearer {resolved_api_key}", "Content-Type": "application/json"} + + +def tokenhub_image_params(params: Mapping[str, object], model: str, drop_params: bool) -> Mapping[str, object]: + n: Final = params.get("n") + if n is not None and n != 1 and not drop_params: + raise UnsupportedParamsError( + status_code=400, + message=f"Tencent TokenHub returns one image per request, so n={n} is not supported for {model}. " + + "To send n=1 instead, set `litellm.drop_params = True`.", + ) + return { + key: value + for key, value in params.items() + if key != "n" and value is not None and not (key == "background" and value == "auto") + } + + +def tokenhub_image_response(raw_response: httpx.Response) -> ImageResponse | None: + try: + parsed: Final = _TokenHubImageResponse.model_validate_json(raw_response.content) + except ValidationError: + return None + total_tokens: Final = parsed.tokenhub_usage.total_tokens if parsed.tokenhub_usage is not None else 0 + first_image: Final = parsed.data[0] if parsed.data else _TokenHubImage() + return ImageResponse( + created=parsed.created, + data=[ImageObject(url=image.url, b64_json=image.b64_json) for image in parsed.data], + size=first_image.size, + output_format=first_image.output_format, + usage=ImageUsage( + input_tokens=0, + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0), + output_tokens=total_tokens, + total_tokens=total_tokens, + ), + ) + + +class TencentImageGenerationConfig(BaseImageGenerationConfig): + @override + def get_supported_openai_params( + self, model: str + ) -> list[OpenAIImageGenerationOptionalParams]: # mutable-ok: mirrors override contract + return list(TOKENHUB_IMAGE_PARAMS) + + @override + def map_openai_params( + self, + non_default_params: Mapping[str, object], + optional_params: Mapping[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: mirrors override contract + return {**optional_params, **tokenhub_image_params(non_default_params, model, drop_params)} + + @override + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + stream: bool | None = None, + ) -> str: + return tokenhub_image_url(api_base) + + @override + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: mirrors override contract + return dict(tokenhub_image_headers(headers, api_key, model)) + + @override + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + headers: Mapping[str, str], + ) -> dict[str, object]: # mutable-ok: mirrors override contract + extra_body: Final = _JSON_OBJECT.validate_python(optional_params.get("extra_body") or {}) + params: Final = {key: value for key, value in optional_params.items() if key != "extra_body"} + return {"model": model, "prompt": prompt, **params, **extra_body} + + @override + def transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: Mapping[str, object], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + encoding: object, + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ImageResponse: + image_response: Final = tokenhub_image_response(raw_response) + if image_response is None: + raise BaseLLMException( + status_code=raw_response.status_code, + message=raw_response.text, + headers=raw_response.headers, + ) + return image_response diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 9f592398f74..0339533e1f5 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) @@ -357,7 +358,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): ) -> VertexAgentEngineResponseIterator: """Return a streaming iterator for SSE responses.""" return VertexAgentEngineResponseIterator( - streaming_response=raw_response.iter_lines(), + streaming_response=iter_sse_lines(raw_response.iter_text()), sync_stream=True, ) @@ -465,7 +466,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): # Create iterator for SSE stream (async) completion_stream: Final = VertexAgentEngineResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 9acbc9c459c..c506192f253 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -41,6 +41,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines, iter_sse_lines from litellm.types.llms.anthropic import AnthropicThinkingParam from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.gemini import BidiGenerateContentServerMessage @@ -2743,10 +2744,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): chunk_indices = support.get("groundingChunkIndices", []) if start_index is not None and end_index is not None and chunk_indices: - # Use the first chunk's URL for the annotation - first_chunk_idx = chunk_indices[0] - if first_chunk_idx in chunk_to_uri_map: - uri_info = chunk_to_uri_map[first_chunk_idx] + for chunk_index in chunk_indices: + if chunk_index not in chunk_to_uri_map: + continue + uri_info = chunk_to_uri_map[chunk_index] url_citation: ChatCompletionAnnotationURLCitation = { "start_index": start_index, @@ -3241,7 +3242,7 @@ async def make_call( raise completion_stream: Final = ModelResponseIterator( - streaming_response=response.aiter_lines(), + streaming_response=aiter_sse_lines(response.aiter_text()), sync_stream=False, logging_obj=logging_obj, response_headers=response.headers, @@ -3312,7 +3313,7 @@ def make_sync_call( raise completion_stream: Final = ModelResponseIterator( - streaming_response=response.iter_lines(), + streaming_response=iter_sse_lines(response.iter_text()), sync_stream=True, logging_obj=logging_obj, response_headers=response.headers, @@ -4190,7 +4191,7 @@ class ModelResponseIterator: log_failed_hop(e) return self._truncated_stream_end(token) self._start_hop(hop_response, token) - self.streaming_response = hop_response.aiter_lines() + self.streaming_response = aiter_sse_lines(hop_response.aiter_text()) self.async_response_iterator = self.streaming_response return None @@ -4204,7 +4205,7 @@ class ModelResponseIterator: log_failed_hop(e) return self._truncated_stream_end(token) self._start_hop(hop_response, token) - self.streaming_response = hop_response.iter_lines() + self.streaming_response = iter_sse_lines(hop_response.iter_text()) self.response_iterator = self.streaming_response return None diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 9fa86ec30f5..ec2fd16f93e 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -9,10 +9,11 @@ The actual message transformation reuses OpenAIGPTConfig since Gemma uses OpenAI """ from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Final, cast import httpx -from pydantic import ValidationError +from pydantic import ConfigDict, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypeAliasType, TypedDict from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import ( @@ -32,6 +33,17 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +class _PredictionsEnvelope(TypedDict): + predictions: NotRequired[ReadOnly[dict[str, object]]] + + +_PredictionsEnvelopeType = TypeAliasType("_PredictionsEnvelopeType", _PredictionsEnvelope) +_PREDICTIONS_ENVELOPE: Final[TypeAdapter[_PredictionsEnvelope]] = TypeAdapter( + _PredictionsEnvelopeType, + config=ConfigDict(strict=True, extra="allow"), +) + + def parse_vertex_gemma_container_error(predictions: object) -> VertexGemmaContainerError | None: try: return VertexGemmaContainerError.model_validate(predictions) @@ -142,7 +154,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): def _unwrap_predictions_response( self, - response_json: dict[str, Any], + response_json: _PredictionsEnvelope, ) -> dict[str, object]: """ Unwrap the Vertex Gemma predictions format to OpenAI format. @@ -354,7 +366,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): message=f"Request failed: {response.text}", ) - response_json: Final = response.json() + response_json: Final = _PREDICTIONS_ENVELOPE.validate_python(response.json()) # Unwrap predictions to get OpenAI-compatible response openai_response: Final = self._unwrap_predictions_response(response_json) @@ -449,7 +461,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): message=f"Request failed: {response.text}", ) - response_json: Final = response.json() + response_json: Final = _PREDICTIONS_ENVELOPE.validate_python(response.json()) # Unwrap predictions to get OpenAI-compatible response openai_response: Final = self._unwrap_predictions_response(response_json) diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 747ea21e4fb..4076ed06fc2 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -3,6 +3,7 @@ from types import UnionType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, Union, get_args, get_origin import httpx +from pydantic import ConfigDict, TypeAdapter from pydantic import fields as pyd_fields import litellm @@ -37,6 +38,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_RESPONSE_BODY: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class _EventModelClass(Protocol): @property @@ -420,7 +423,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): @staticmethod def _parsed_response_body(raw_response: httpx.Response) -> dict[str, object]: - return raw_response.json() + return _RESPONSE_BODY.validate_python(raw_response.json()) @staticmethod def _annotation_origin(annotation: object) -> object: diff --git a/litellm/main.py b/litellm/main.py index 75bd4431120..41c4302499e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6124,7 +6124,18 @@ def completion( optional_params, rs_val = strip_reasoning_summary_aliases_from_optional_params(optional_params) - if isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort: + if custom_llm_provider == "openai" and isinstance(resolved_reasoning_effort, dict): + from litellm.responses.utils import openai_bridge_reasoning + + optional_params["reasoning_effort"] = openai_bridge_reasoning( + requested=resolved_reasoning_effort, + effort_kept="reasoning_effort" in optional_params, + summary_alias=rs_val, + model=model, + drop_params=drop_unsupported_reasoning, + additional_drop_params=routing_kwargs.get("additional_drop_params"), + ) + elif isinstance(resolved_reasoning_effort, dict) and "summary" in resolved_reasoning_effort: optional_params["reasoning_effort"] = resolved_reasoning_effort elif rs_val is not None: eff: Final = optional_params.get("reasoning_effort", resolved_reasoning_effort) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1b65c87c315..8719565a96f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -21009,10 +21009,11 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-fable-5": { - "cache_creation_input_token_cost": 1.250004e-05, + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_1hr": 1.999998e-05, "cache_read_input_token_cost": 1.00002e-06, - "input_cost_per_token": 1.000006e-05, - "input_dbu_cost_per_token": 0.000142858, + "input_cost_per_token": 9.99999e-06, + "input_dbu_cost_per_token": 0.00014286, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -21021,7 +21022,7 @@ }, "mode": "chat", "output_cost_per_token": 5.000002e-05, - "output_dbu_cost_per_token": 0.000714286, + "output_dbu_cost_per_token": 0.00071429, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21037,10 +21038,11 @@ "thinking_always_on": true }, "databricks/databricks-claude-fable-5-1": { - "cache_creation_input_token_cost": 1.250004e-05, - "cache_read_input_token_cost": 2.5004e-07, - "input_cost_per_token": 1.000006e-05, - "input_dbu_cost_per_token": 0.000142858, + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_1hr": 1.999998e-05, + "cache_read_input_token_cost": 2.4997e-07, + "input_cost_per_token": 9.99999e-06, + "input_dbu_cost_per_token": 0.00014286, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -21049,7 +21051,7 @@ }, "mode": "chat", "output_cost_per_token": 5.000002e-05, - "output_dbu_cost_per_token": 0.000714286, + "output_dbu_cost_per_token": 0.00071429, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21067,6 +21069,7 @@ }, "databricks/databricks-claude-haiku-4-5": { "cache_creation_input_token_cost": 1.24999e-06, + "cache_creation_input_token_cost_above_1hr": 1.99997e-06, "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, @@ -21090,9 +21093,10 @@ }, "databricks/databricks-claude-opus-4": { "cache_creation_input_token_cost": 1.874999e-05, + "cache_creation_input_token_cost_above_1hr": 2.999997e-05, "cache_read_input_token_cost": 1.50003e-06, - "input_cost_per_token": 1.5000020000000002e-05, - "input_dbu_cost_per_token": 0.000214286, + "input_cost_per_token": 1.500002e-05, + "input_dbu_cost_per_token": 0.00021429, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -21100,8 +21104,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 7.500003000000001e-05, - "output_dbu_cost_per_token": 0.001071429, + "output_cost_per_token": 7.50001e-05, + "output_dbu_cost_per_token": 0.0010714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21113,9 +21117,10 @@ }, "databricks/databricks-claude-opus-4-1": { "cache_creation_input_token_cost": 1.874999e-05, + "cache_creation_input_token_cost_above_1hr": 2.999997e-05, "cache_read_input_token_cost": 1.50003e-06, - "input_cost_per_token": 1.5000020000000002e-05, - "input_dbu_cost_per_token": 0.000214286, + "input_cost_per_token": 1.500002e-05, + "input_dbu_cost_per_token": 0.00021429, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -21123,8 +21128,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 7.500003000000001e-05, - "output_dbu_cost_per_token": 0.001071429, + "output_cost_per_token": 7.50001e-05, + "output_dbu_cost_per_token": 0.0010714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21136,6 +21141,7 @@ }, "databricks/databricks-claude-opus-4-5": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21146,8 +21152,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.00035714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21160,6 +21166,7 @@ }, "databricks/databricks-claude-opus-4-6": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21170,8 +21177,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.00035714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": false, "supports_function_calling": true, @@ -21184,6 +21191,7 @@ }, "databricks/databricks-claude-opus-4-7": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21195,7 +21203,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 2048, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21210,6 +21218,7 @@ }, "databricks/databricks-claude-opus-4-8": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21221,7 +21230,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21237,6 +21246,7 @@ }, "databricks/databricks-claude-opus-5": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21248,7 +21258,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21344,8 +21354,9 @@ }, "databricks/databricks-claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.74997e-06, + "cache_creation_input_token_cost_above_1hr": 5.99998e-06, "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 200000, @@ -21354,8 +21365,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.00021429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21367,8 +21378,9 @@ }, "databricks/databricks-claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.74997e-06, + "cache_creation_input_token_cost_above_1hr": 5.99998e-06, "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, @@ -21377,8 +21389,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.00021429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": false, "supports_function_calling": true, @@ -21390,19 +21402,20 @@ "prompt_cache_min_tokens": 1024 }, "databricks/databricks-claude-sonnet-5": { - "cache_creation_input_token_cost": 3.74997e-06, - "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.99999e-06, - "input_dbu_cost_per_token": 4.2857e-05, + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_1hr": 4.00001e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, "metadata": { - "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. Introductory launch rates of 28.571 input / 142.857 output / 35.714 cache write / 2.857 cache read DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 1.500002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21546,7 +21559,10 @@ }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_audio_token_cost": 6.251e-08, "cache_read_input_token_cost": 3.122e-08, + "input_cost_per_audio_token": 6.2503e-07, + "input_cost_per_image_token": 3.1248e-07, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -21573,16 +21589,28 @@ ], "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_video_input": true, + "input_cost_per_token_priority": 5.6252e-07, + "output_cost_per_token_priority": 3.37498e-06, + "cache_read_input_token_cost_priority": 5.628e-08, + "cache_creation_input_token_cost_priority": 5.6252e-07, + "input_cost_per_audio_token_priority": 1.12497e-06, + "cache_read_input_audio_token_cost_priority": 1.1249e-07 }, "databricks/databricks-gemini-3-1-flash-image": { + "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_token_cost": 6.2503e-07, + "input_cost_per_image_token": 6.2503e-07, + "input_cost_per_token": 6.2503e-07, "litellm_provider": "databricks", "max_input_tokens": 131072, "max_output_tokens": 32768, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_image_token": 7.50001e-05, + "output_cost_per_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21596,13 +21624,19 @@ "supports_vision": true }, "databricks/databricks-gemini-3-pro-image": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-06, + "input_cost_per_image_token": 2.49998e-06, + "input_cost_per_token": 2.49998e-06, "litellm_provider": "databricks", "max_input_tokens": 65536, "max_output_tokens": 32768, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_image_token": 0.0001500002, + "output_cost_per_token": 1.500002e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21618,7 +21652,9 @@ "databricks/databricks-gemini-3-1-pro": { "cache_creation_input_token_cost": 2.49998e-06, "cache_read_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost_above_200k_tokens": 5.0001e-07, "input_cost_per_token": 2.49998e-06, + "input_cost_per_token_above_200k_tokens": 5.00003e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", "max_input_tokens": 1048576, @@ -21628,12 +21664,9 @@ }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token_above_200k_tokens": 2.25e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_tool_choice": true, - "supports_reasoning": true, "supported_modalities": [ "text", "image", @@ -21643,13 +21676,27 @@ "supported_output_modalities": [ "text" ], - "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "input_cost_per_token_priority": 4.50002e-06, + "output_cost_per_token_priority": 2.699998e-05, + "cache_read_input_token_cost_priority": 4.5003e-07, + "cache_creation_input_token_cost_priority": 4.50002e-06, + "input_cost_per_token_above_200k_tokens_priority": 8.99997e-06, + "output_cost_per_token_above_200k_tokens_priority": 4.049997e-05, + "cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07 }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_audio_token_cost": 1.2502e-07, "cache_read_input_token_cost": 6.251e-08, + "input_cost_per_audio_token": 1.24999e-06, + "input_cost_per_image_token": 6.2503e-07, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -21676,7 +21723,13 @@ ], "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_video_input": true, + "input_cost_per_token_priority": 1.12497e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.1249e-07, + "cache_creation_input_token_cost_priority": 1.12497e-06, + "input_cost_per_audio_token_priority": 2.25001e-06, + "cache_read_input_audio_token_cost_priority": 2.2498e-07 }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21698,58 +21751,9 @@ "supports_tool_choice": true }, "databricks/databricks-gemini-3-8-flash": { - "litellm_provider": "databricks", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." - }, - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_function_calling": true, - "supports_reasoning": true, - "supports_sampling_params": false, - "supports_tool_choice": true, - "supports_vision": true - }, - "databricks/databricks-gemini-3-7-flash": { - "litellm_provider": "databricks", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." - }, - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_function_calling": true, - "supports_reasoning": true, - "supports_sampling_params": false, - "supports_tool_choice": true, - "supports_vision": true - }, - "databricks/databricks-gemini-3-6-flash": { - "cache_creation_input_token_cost": 1.87502e-06, - "cache_read_input_token_cost": 1.8753e-07, - "input_cost_per_token": 1.87502e-06, - "input_dbu_cost_per_token": 2.6786e-05, + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -21757,8 +21761,73 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 9.37503e-06, - "output_dbu_cost_per_token": 0.000133929, + "output_cost_per_token": 3.74997e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 + }, + "databricks/databricks-gemini-3-7-flash": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 3.74997e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 + }, + "databricks/databricks-gemini-3-6-flash": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, + "input_dbu_cost_per_token": 1.0714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 3.74997e-06, + "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21774,7 +21843,11 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 }, "databricks/databricks-gemini-3-5-flash": { "cache_creation_input_token_cost": 1.87502e-06, @@ -21804,7 +21877,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 3.37498e-06, + "output_cost_per_token_priority": 2.025002e-05, + "cache_read_input_token_cost_priority": 3.3747e-07, + "cache_creation_input_token_cost_priority": 3.37498e-06 }, "databricks/databricks-gemini-3-5-flash-lite": { "cache_creation_input_token_cost": 3.7499e-07, @@ -21836,7 +21913,11 @@ "supports_response_schema": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 6.7501e-07, + "output_cost_per_token_priority": 5.62499e-06, + "cache_read_input_token_cost_priority": 6.748e-08, + "cache_creation_input_token_cost_priority": 6.7501e-07 }, "databricks/databricks-gemma-3-12b": { "cache_creation_input_token_cost": 1.5001e-07, @@ -21879,7 +21960,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "input_cost_per_token_priority": 2.45e-06, + "output_cost_per_token_priority": 7.7e-06, + "cache_read_input_token_cost_priority": 4.55e-07, + "cache_creation_input_token_cost_priority": 2.45e-06 }, "databricks/databricks-glm-5-3": { "cache_creation_input_token_cost": 1.4e-06, @@ -21909,7 +21994,11 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 2.45e-06, + "output_cost_per_token_priority": 7.7e-06, + "cache_read_input_token_cost_priority": 4.55e-07, + "cache_creation_input_token_cost_priority": 2.45e-06 }, "databricks/databricks-glm-5-3-flash": { "cache_creation_input_token_cost": 1.5001e-07, @@ -21938,7 +22027,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 2.625e-07, + "output_cost_per_token_priority": 8.75e-07, + "cache_read_input_token_cost_priority": 5.25e-08, + "cache_creation_input_token_cost_priority": 2.625e-07 }, "databricks/databricks-gpt-5": { "cache_creation_input_token_cost": 1.24999e-06, @@ -21957,7 +22050,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 2.49998e-06, + "output_cost_per_token_priority": 1.999998e-05, + "cache_read_input_token_cost_priority": 2.4997e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06 }, "databricks/databricks-gpt-5-1": { "cache_creation_input_token_cost": 1.24999e-06, @@ -21976,7 +22073,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 2.49998e-06, + "output_cost_per_token_priority": 1.999998e-05, + "cache_read_input_token_cost_priority": 2.4997e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06 }, "databricks/databricks-gpt-5-2": { "cache_creation_input_token_cost": 1.75e-06, @@ -21995,7 +22096,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 3.5e-06, + "output_cost_per_token_priority": 2.8e-05, + "cache_read_input_token_cost_priority": 3.5e-07, + "cache_creation_input_token_cost_priority": 3.5e-06 }, "databricks/databricks-gpt-5-3-codex": { "cache_creation_input_token_cost": 1.75e-06, @@ -22014,12 +22119,18 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 3.5e-06, + "output_cost_per_token_priority": 2.8e-05, + "cache_read_input_token_cost_priority": 3.5e-07, + "cache_creation_input_token_cost_priority": 3.5e-06 }, "databricks/databricks-gpt-5-4": { "cache_creation_input_token_cost": 2.49998e-06, "cache_read_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.0001e-07, "input_cost_per_token": 2.49998e-06, + "input_cost_per_token_above_272k_tokens": 5.00003e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22029,6 +22140,7 @@ }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22042,7 +22154,14 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.00003e-06, + "output_cost_per_token_priority": 2.999997e-05, + "cache_read_input_token_cost_priority": 5.0001e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 5.00003e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.999997e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07 }, "databricks/databricks-gpt-5-4-mini": { "cache_creation_input_token_cost": 7.4998e-07, @@ -22070,7 +22189,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.50003e-06, + "output_cost_per_token_priority": 8.99997e-06, + "cache_read_input_token_cost_priority": 1.5001e-07, + "cache_creation_input_token_cost_priority": 1.50003e-06 }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, @@ -22102,8 +22225,11 @@ }, "databricks/databricks-gpt-5-6-sol": { "cache_creation_input_token_cost": 5.00003e-06, + "cache_creation_input_token_cost_above_272k_tokens": 9.99999e-06, "cache_read_input_token_cost": 3.9998e-07, + "cache_read_input_token_cost_above_272k_tokens": 8.0003e-07, "input_cost_per_token": 4.00001e-06, + "input_cost_per_token_above_272k_tokens": 8.00002e-06, "input_dbu_cost_per_token": 5.7143e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22113,6 +22239,7 @@ }, "mode": "chat", "output_cost_per_token": 1.999998e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, "output_dbu_cost_per_token": 0.000285714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22126,12 +22253,23 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 8.00002e-06, + "output_cost_per_token_priority": 4.000003e-05, + "cache_read_input_token_cost_priority": 8.0003e-07, + "cache_creation_input_token_cost_priority": 9.99999e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.599997e-05, + "output_cost_per_token_above_272k_tokens_priority": 6.000001e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05 }, "databricks/databricks-gpt-5-6-terra": { "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, "cache_read_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost_above_272k_tokens": 3.9998e-07, "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22141,6 +22279,7 @@ }, "mode": "chat", "output_cost_per_token": 1.200003e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, "output_dbu_cost_per_token": 0.000171429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22154,12 +22293,23 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 2.399999e-05, + "cache_read_input_token_cost_priority": 3.9998e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.600002e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, "databricks/databricks-gpt-5-6-luna": { "cache_creation_input_token_cost": 2.4997e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.0001e-07, "cache_read_input_token_cost": 2.002e-08, + "cache_read_input_token_cost_above_272k_tokens": 3.997e-08, "input_cost_per_token": 1.9999e-07, + "input_cost_per_token_above_272k_tokens": 3.9998e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22169,6 +22319,7 @@ }, "mode": "chat", "output_cost_per_token": 1.20001e-06, + "output_cost_per_token_above_272k_tokens": 1.79998e-06, "output_dbu_cost_per_token": 1.7143e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22182,12 +22333,22 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 3.9998e-07, + "output_cost_per_token_priority": 2.40002e-06, + "cache_read_input_token_cost_priority": 3.997e-08, + "cache_creation_input_token_cost_priority": 5.0001e-07, + "input_cost_per_token_above_272k_tokens_priority": 8.0003e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.60003e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06 }, "databricks/databricks-gpt-5-5": { "cache_creation_input_token_cost": 5.00003e-06, "cache_read_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.00002e-06, "input_cost_per_token": 5.00003e-06, + "input_cost_per_token_above_272k_tokens": 9.99999e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22197,6 +22358,7 @@ }, "mode": "responses", "output_cost_per_token": 2.999997e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_dbu_cost_per_token": 0.000428571, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22210,13 +22372,21 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.249997e-05, + "output_cost_per_token_priority": 7.50001e-05, + "cache_read_input_token_cost_priority": 1.24999e-06, + "cache_creation_input_token_cost_priority": 1.249997e-05, + "input_cost_per_token_above_272k_tokens_priority": 1.249997e-05, + "output_cost_per_token_above_272k_tokens_priority": 7.50001e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06 }, "databricks/databricks-gpt-5-5-pro": { "cache_creation_input_token_cost": 2.999997e-05, "cache_read_input_token_cost": 2.999997e-05, "input_cost_per_token": 2.999997e-05, - "input_dbu_cost_per_token": 0.000428571, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_dbu_cost_per_token": 0.00042857, "litellm_provider": "databricks", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -22224,8 +22394,9 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "responses", - "output_cost_per_token": 0.00018000003, - "output_dbu_cost_per_token": 0.002571429, + "output_cost_per_token": 0.0001800001, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_dbu_cost_per_token": 0.0025714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -22257,7 +22428,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 4.5003e-07, + "output_cost_per_token_priority": 3.60003e-06, + "cache_read_input_token_cost_priority": 4.501e-08, + "cache_creation_input_token_cost_priority": 4.5003e-07 }, "databricks/databricks-gpt-5-nano": { "cache_creation_input_token_cost": 4.998e-08, @@ -22407,7 +22582,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.25e-06, + "output_cost_per_token_priority": 2.625e-05, + "cache_read_input_token_cost_priority": 5.25e-07, + "cache_creation_input_token_cost_priority": 5.25e-06 }, "databricks/system.ai.kimi-k3": { "litellm_provider": "databricks", @@ -22489,7 +22668,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": false, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 3.85e-07, + "output_cost_per_token_priority": 3.85e-06, + "cache_read_input_token_cost_priority": 3.85e-07, + "cache_creation_input_token_cost_priority": 3.85e-07 }, "databricks/databricks-qwen3-next-80b-a3b-instruct": { "cache_creation_input_token_cost": 1.5001e-07, @@ -22544,6 +22727,25 @@ "output_dbu_cost_per_token": 0.0, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" }, + "databricks/databricks-typesafeai-jev-1-14-0": { + "cache_creation_input_token_cost": 4.2e-08, + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 4.2e-08, + "input_dbu_cost_per_token": 6e-07, + "litellm_provider": "databricks", + "metadata": { + "notes": "TypeSafe AI Jev 1.14.0, provider-hosted by TypeSafe AI and served through Unity Gateway, queried through the System One API. The proprietary pricing page lists 'Jev 1.14' at 0.600 DBU per 1M input tokens; input_cost_per_token is that times $0.070 per DBU. Output is free. No cache rates are published, so cache tokens bill at the input rate." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "output_dbu_cost_per_token": 0.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supports_vision": true + }, "databricks/ai_decide": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, @@ -28651,6 +28853,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -28708,6 +28911,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -28761,12 +28965,13 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -28812,6 +29017,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -28857,9 +29063,10 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -28904,6 +29111,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -29470,10 +29678,12 @@ }, "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, @@ -29488,6 +29698,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -29499,6 +29710,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -30027,6 +30239,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -30038,6 +30251,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -35598,6 +35812,13 @@ "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35675,6 +35896,13 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35851,6 +36079,13 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35902,6 +36137,13 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36017,6 +36259,13 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36066,6 +36315,13 @@ "mode": "responses", "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36107,6 +36363,13 @@ "mode": "responses", "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36194,6 +36457,15 @@ "output_cost_per_token_flex": 2.5e-05, "output_cost_per_token_priority": 0.0001, "output_cost_per_token_ultrafast": 0.0003, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36278,6 +36550,15 @@ "output_cost_per_token_above_272k_tokens_batches": 7.5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36370,6 +36651,15 @@ "output_cost_per_token_above_272k_tokens_batches": 3.75e-07, "output_cost_per_token_flex": 2.5e-07, "output_cost_per_token_priority": 1e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36453,6 +36743,15 @@ "output_cost_per_token_above_272k_tokens_batches": 1.5e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36534,6 +36833,15 @@ "output_cost_per_token_above_272k_tokens_batches": 1.5e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36617,6 +36925,15 @@ "output_cost_per_token_above_272k_tokens_batches": 9e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36699,6 +37016,15 @@ "output_cost_per_token_above_272k_tokens_batches": 9e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36936,6 +37262,14 @@ "reasoning_effort_levels": [ "medium" ], + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -36983,6 +37317,14 @@ "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37045,6 +37387,14 @@ "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37099,6 +37449,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37150,6 +37508,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37209,6 +37575,14 @@ "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_priority": 3e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -37266,6 +37640,14 @@ "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_priority": 3e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -37315,6 +37697,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37367,6 +37757,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37421,6 +37819,14 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37475,6 +37881,14 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37526,6 +37940,14 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37578,6 +38000,14 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37624,6 +38054,13 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37666,6 +38103,13 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37716,6 +38160,13 @@ "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37826,6 +38277,13 @@ "mode": "responses", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37908,6 +38366,13 @@ "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37959,6 +38424,13 @@ "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38007,6 +38479,13 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38055,6 +38534,13 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -41845,6 +42331,22 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/deepseek-v4-1-flash-1017": { + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.mistral.ai/models/deepseek-v4-flash-1017", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/zai-glm-5-3": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -43805,6 +44307,13 @@ "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -43913,6 +44422,13 @@ "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -43992,6 +44508,13 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -44029,6 +44552,13 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -51392,6 +51922,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -51450,6 +51981,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -51484,13 +52016,14 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, - "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -51516,6 +52049,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -51543,10 +52077,11 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "supports_reasoning": false, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -51571,6 +52106,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -61468,6 +62004,39 @@ "supports_reasoning": true, "supports_vision": false }, + "tencent/seedream-image-v5.0-flash": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.018, + "output_cost_per_image_token": 1.6e-06, + "source": "https://cloud.tencent.com/document/product/1823/130055", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "tencent/seedream-image-v5.0-lite": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.035, + "output_cost_per_image_token": 1.6e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "tencent/seedream-image-v5.0-pro": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.09, + "output_cost_per_image_token": 1.6e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, "cognition/swe-1.6": { "input_cost_per_token": 5e-07, "output_cost_per_token": 2.5e-06, @@ -65847,6 +66416,180 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "zerogpu/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/glm-5.3-flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1048576, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-4.1-mini": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 4e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-5.6-luna": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-6-luna": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "zerogpu/gpt-oss-120b": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/LFM2.5-1.2B-Instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-08, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/LFM2.5-1.2B-Thinking": { + "input_cost_per_token": 2e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-08, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/llama-3.1-8b-instruct-fast": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/llama-3.2-3b-instruct-fast": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 5e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/qwen3-30b-a3b-fp8": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4.5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/qwen3.6-35b-a3b": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "together_ai/Qwen/Qwen3.8-Flash": { "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", @@ -73272,6 +74015,15 @@ "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "output_cost_per_token_ultrafast": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -73666,6 +74418,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, @@ -73711,6 +74464,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, @@ -76771,6 +77525,15 @@ "output_cost_per_token_batches": 2.5e-06, "output_cost_per_token_flex": 2.5e-06, "output_cost_per_token_priority": 1e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, @@ -77095,6 +77858,577 @@ "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, + "databricks/databricks-claude-haiku-5-5": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_creation_input_token_cost_above_100k_tokens": 6.2503e-07, + "cache_creation_input_token_cost_above_1hr": 1.9999e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.00002e-06, + "cache_read_input_token_cost": 1.001e-08, + "cache_read_input_token_cost_above_100k_tokens": 4.998e-08, + "input_cost_per_token": 1.0003e-07, + "input_cost_per_token_above_100k_tokens": 5.0001e-07, + "input_dbu_cost_per_token": 1.429e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_cost_per_token_above_100k_tokens": 2.49998e-06, + "output_dbu_cost_per_token": 7.143e-06, + "prompt_cache_min_tokens": 512.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_anthropic_thinking_payload": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "thinking_always_on": false + }, + "databricks/databricks-claude-sonnet-5-5": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_1hr": 4.00001e-06, + "cache_read_input_token_cost": 1.0003e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, + "prompt_cache_min_tokens": 512.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_anthropic_thinking_payload": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": false, + "thinking_always_on": true + }, + "databricks/databricks-deepseek-v4-1-flash": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 3.003e-08, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 3.0002e-07, + "input_dbu_cost_per_token": 4.286e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. The docs list only databricks-deepseek-v4-1-flash (retiring 2026-11-19) and databricks-deepseek-v4-1-flash-1017, so the pricing page's 'DeepSeek V4.1 Flash (0910)' row prices the unsuffixed endpoint." + }, + "mode": "chat", + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, + "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 5.25e-07, + "output_cost_per_token_priority": 2.1e-06, + "cache_read_input_token_cost_priority": 5.25e-08, + "cache_creation_input_token_cost_priority": 5.25e-07 + }, + "databricks/databricks-deepseek-v4-1-flash-1017": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 3.0002e-07, + "input_dbu_cost_per_token": 4.286e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gemini-2-5-flash-lite": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_read_input_audio_token_cost": 3.752e-08, + "cache_read_input_token_cost": 1.253e-08, + "input_cost_per_audio_token": 3.7499e-07, + "input_cost_per_image_token": 1.2502e-07, + "input_cost_per_token": 1.2502e-07, + "input_dbu_cost_per_token": 1.786e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "databricks/databricks-gemini-3-1-flash-lite-image": { + "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_token_cost": 3.1248e-07, + "input_cost_per_image_token": 3.1248e-07, + "input_cost_per_token": 3.1248e-07, + "input_dbu_cost_per_token": 4.464e-06, + "litellm_provider": "databricks", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_image_token": 3.749998e-05, + "output_cost_per_token": 1.87502e-06, + "output_dbu_cost_per_token": 2.6786e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_vision": true + }, + "databricks/databricks-gemini-4-argon": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 1.2502e-07, + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 1.249997e-05, + "output_dbu_cost_per_token": 0.00017857, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "input_cost_per_token_priority": 4.50002e-06, + "output_cost_per_token_priority": 2.250003e-05, + "cache_read_input_token_cost_priority": 2.2498e-07, + "cache_creation_input_token_cost_priority": 4.50002e-06 + }, + "databricks/databricks-gpt-5-4-pro": { + "cache_creation_input_token_cost": 2.999997e-05, + "cache_read_input_token_cost": 2.999997e-05, + "input_cost_per_token": 2.999997e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_dbu_cost_per_token": 0.00042857, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 0.0001800001, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_dbu_cost_per_token": 0.0025714, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-1-sol": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, + "cache_read_input_token_cost": 1.0003e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 9.99999e-06, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 1.99997e-05, + "cache_read_input_token_cost_priority": 1.9999e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.99999e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 + }, + "databricks/databricks-gpt-6-1-terra": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.49998e-06, + "cache_read_input_token_cost": 4.998e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.0003e-07, + "input_cost_per_token": 1.00002e-06, + "input_cost_per_token_above_272k_tokens": 1.99997e-06, + "input_dbu_cost_per_token": 1.4286e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.00003e-06, + "output_cost_per_token_above_272k_tokens": 7.50001e-06, + "output_dbu_cost_per_token": 7.1429e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.99997e-06, + "output_cost_per_token_priority": 9.99999e-06, + "cache_read_input_token_cost_priority": 1.0003e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06, + "input_cost_per_token_above_272k_tokens_priority": 4.00001e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.500002e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06 + }, + "databricks/databricks-gpt-6-astra": { + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.99997e-06, + "input_cost_per_token": 9.99999e-06, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_dbu_cost_per_token": 0.00014286, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.000002e-05, + "output_cost_per_token_above_272k_tokens": 7.50001e-05, + "output_dbu_cost_per_token": 0.00071429, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.999998e-05, + "output_cost_per_token_priority": 9.99999e-05, + "cache_read_input_token_cost_priority": 1.99997e-06, + "cache_creation_input_token_cost_priority": 2.500001e-05, + "input_cost_per_token_above_272k_tokens_priority": 4.000003e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.0001500002, + "cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05 + }, + "databricks/databricks-gpt-6-luna": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_creation_input_token_cost_above_272k_tokens": 2.4997e-07, + "cache_read_input_token_cost": 1.001e-08, + "cache_read_input_token_cost_above_272k_tokens": 2.002e-08, + "input_cost_per_token": 1.0003e-07, + "input_cost_per_token_above_272k_tokens": 1.9999e-07, + "input_dbu_cost_per_token": 1.429e-06, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 5.0001e-07, + "output_cost_per_token_above_272k_tokens": 7.497e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.9999e-07, + "output_cost_per_token_priority": 1.0003e-06, + "cache_read_input_token_cost_priority": 2.002e-08, + "cache_creation_input_token_cost_priority": 2.4997e-07, + "input_cost_per_token_above_272k_tokens_priority": 3.9998e-07, + "output_cost_per_token_above_272k_tokens_priority": 1.5001e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07 + }, + "databricks/databricks-gpt-6-sol": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, + "cache_read_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost_above_272k_tokens": 3.9998e-07, + "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 9.99999e-06, + "output_cost_per_token_above_272k_tokens": 1.50003e-05, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 1.99997e-05, + "cache_read_input_token_cost_priority": 3.9998e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.99999e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 + }, + "databricks/databricks-gpt-image-1": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 9.99999e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 4.000003e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-image-1-5": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 8.00002e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 3.200001e-05, + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_pdf_input": true, + "supports_vision": true + }, + "databricks/databricks-gpt-image-1-mini": { + "cache_creation_input_token_cost": 1.99997e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_image_token": 2.49998e-06, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 8.00002e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-image-2": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 8.00002e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 2.999997e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_pdf_input": true, + "supports_vision": true + }, + "databricks/databricks-grok-4-7": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 6.2503e-07, + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 7.50001e-06, + "output_dbu_cost_per_token": 0.00010714, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-kimi-k2-7": { + "cache_creation_input_token_cost": 9.4997e-07, + "cache_read_input_token_cost": 1.8998e-07, + "input_cost_per_token": 9.4997e-07, + "input_dbu_cost_per_token": 1.3571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 4.00001e-06, + "output_dbu_cost_per_token": 5.7143e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { "cache_read_input_token_cost": 2.1e-08, "deprecation_date": "2026-08-14", diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 1ab03fb3fdc..0a6c2bcb2a9 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1225,7 +1225,7 @@ "responses": true, "text_completion": true, "embeddings": true, - "image_generations": false, + "image_generations": true, "image_edits": false, "audio_transcriptions": false, "audio_speech": false, @@ -3793,8 +3793,8 @@ "responses": true, "text_completion": true, "embeddings": false, - "image_generations": false, - "image_edits": false, + "image_generations": true, + "image_edits": true, "audio_transcriptions": false, "audio_speech": false, "rerank": false, @@ -4437,6 +4437,32 @@ "generateContent": true } }, + "zerogpu": { + "display_name": "ZeroGPU (`zerogpu`)", + "url": "https://docs.litellm.ai/docs/providers/zerogpu", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } + }, "ragflow": { "display_name": "RAGFlow (`ragflow`)", "url": "https://docs.litellm.ai/docs/providers/ragflow", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 2aaa987ba6a..39adb99a0d0 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3487,6 +3487,13 @@ "title": "Entities", "type": "object" }, + "guardrails": { + "additionalProperties": { + "$ref": "#/components/schemas/GuardrailMetric" + }, + "title": "Guardrails", + "type": "object" + }, "mcp_servers": { "additionalProperties": { "$ref": "#/components/schemas/MetricWithMetadata" @@ -3802,6 +3809,74 @@ "title": "ExportType", "type": "string" }, + "GuardrailKeyMetric": { + "properties": { + "metadata": { + "$ref": "#/components/schemas/KeyMetadata" + }, + "metrics": { + "$ref": "#/components/schemas/GuardrailSpendMetrics" + } + }, + "required": [ + "metrics", + "metadata" + ], + "title": "GuardrailKeyMetric", + "type": "object" + }, + "GuardrailMetadata": { + "properties": { + "guardrail_name": { + "title": "Guardrail Name", + "type": "string" + } + }, + "required": [ + "guardrail_name" + ], + "title": "GuardrailMetadata", + "type": "object" + }, + "GuardrailMetric": { + "properties": { + "api_key_breakdown": { + "additionalProperties": { + "$ref": "#/components/schemas/GuardrailKeyMetric" + }, + "title": "Api Key Breakdown", + "type": "object" + }, + "metadata": { + "$ref": "#/components/schemas/GuardrailMetadata" + }, + "metrics": { + "$ref": "#/components/schemas/GuardrailSpendMetrics" + } + }, + "required": [ + "metrics", + "metadata" + ], + "title": "GuardrailMetric", + "type": "object" + }, + "GuardrailSpendMetrics": { + "properties": { + "spend": { + "default": 0.0, + "title": "Spend", + "type": "number" + }, + "untracked_units": { + "default": 0, + "title": "Untracked Units", + "type": "integer" + } + }, + "title": "GuardrailSpendMetrics", + "type": "object" + }, "HTTPAuthSecurityScheme": { "description": "Defines a security scheme using HTTP authentication.", "properties": { @@ -15015,7 +15090,7 @@ ] }, "post": { - "description": "Create a named custom agent on the Gemini side.\n\nExample:\n```bash\ncurl -X POST \"http://localhost:4000/v1beta/agents\" \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"name\": \"my-custom-slides-agent\",\n \"base_agent\": \"waverunner\",\n \"instructions\": \"You are a helpful assistant that creates slides.\",\n \"base_environment\": {\n \"type\": \"remote\",\n \"sources\": [\n {\"type\": \"gcs\", \"source\": \"gs://eap-templates/slides-skill\",\n \"target\": \"/.agents/skills/slides-skill\"}\n ]\n }\n }'\n```", + "description": "Create a named custom agent on the Gemini side.\n\nExample:\n```bash\ncurl -X POST \"http://localhost:4000/v1beta/agents\" \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"id\": \"my-custom-slides-agent\",\n \"base_agent\": \"antigravity-preview-09-2026\",\n \"system_instruction\": \"You are a helpful assistant that creates slides.\",\n \"agent_config\": {\"type\": \"antigravity\", \"model\": \"gemini-3.8-flash\"},\n \"base_environment\": {\n \"type\": \"remote\",\n \"sources\": [\n {\"type\": \"gcs\", \"source\": \"gs://eap-templates/slides-skill\",\n \"target\": \"/.agents/skills/slides-skill\"}\n ]\n }\n }'\n```", "operationId": "create_gemini_agent_v1beta_agents_post", "responses": { "200": { diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index ff19ba50dd8..c3766c6b9f0 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -27,6 +27,7 @@ from litellm.caching.caching import Cache from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header +from litellm.llms.custom_httpx.sse_lines import aiter_sse_lines from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.a2a.version_convert import ( A2AVersion, @@ -335,7 +336,7 @@ async def _a2a_sse_event_source( "error": {"code": -32603, "message": resp.reason_phrase}, } return - async for line in resp.aiter_lines(): + async for line in aiter_sse_lines(resp.aiter_text()): stripped = line.strip() if not stripped.startswith("data:"): continue diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py index f748a754fe0..32e0f6e8682 100644 --- a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -14,6 +14,8 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.llms.custom_httpx.sse_lines import split_sse_lines_keepends + _MESSAGE_START_EVENT: Final = "message_start" _MESSAGE_START_MARKER: Final = b"message_start" _SSE_DATA_FIELD: Final = "data:" @@ -52,7 +54,7 @@ def _restamped_data_line(line: str, requested_model: str) -> str | None: def _restamped_frame(frame: str, requested_model: str) -> str | None: - lines: Final = frame.splitlines(keepends=True) + lines: Final = split_sse_lines_keepends(frame) restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines) if all(line is None for line in restamped): return None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 99edfb940d7..79ad4824c92 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -370,7 +370,7 @@ class _KeyModelBudgetLimiter(Protocol): class _UserModelBudgetLimiter(Protocol): async def is_user_within_model_budget( - self, user_id: str, user_model_max_budget: Mapping[str, object], model: str + self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None ) -> bool: ... @@ -388,11 +388,15 @@ class _EveryModelBudgetLimiter(Protocol): async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ... async def is_user_within_model_budget( - self, user_id: str, user_model_max_budget: Mapping[str, object], model: str + self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None ) -> bool: ... async def is_end_user_within_model_budget( - self, end_user_id: str, end_user_model_max_budget: Mapping[str, object], model: str + self, + end_user_id: str, + end_user_model_max_budget: Mapping[str, object], + model: str, + team_id: str | None = None, ) -> bool: ... async def is_team_within_model_budget( @@ -503,6 +507,7 @@ async def _check_user_model_budget( user_id=valid_token.user_id, user_model_max_budget=user_model_max_budget, model=model_name, + team_id=valid_token.team_id, ) @@ -547,6 +552,7 @@ async def _is_within_every_model_budget( user_id=valid_token.user_id, user_model_max_budget=user_model_max_budget, model=model, + team_id=valid_token.team_id, ) if ( valid_token.end_user_id is not None @@ -557,6 +563,7 @@ async def _is_within_every_model_budget( end_user_id=valid_token.end_user_id, end_user_model_max_budget=end_user_model_max_budget, model=model, + team_id=valid_token.team_id, ) if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget: await model_max_budget_limiter.is_team_within_model_budget( @@ -3154,6 +3161,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e end_user_id=valid_token.end_user_id, end_user_model_max_budget=end_user_mmb, model=model_name, + team_id=valid_token.team_id, ) # Check 6: Additional Common Checks across jwt + key auth @@ -4733,6 +4741,7 @@ async def _run_post_custom_auth_checks( end_user_id=valid_token.end_user_id, end_user_model_max_budget=end_user_mmb, model=model_name, + team_id=valid_token.team_id, ) # team / user / end_user / project context objects are fetched by diff --git a/litellm/proxy/client/README.md b/litellm/proxy/client/README.md index 1fff68677cc..ceec636567f 100644 --- a/litellm/proxy/client/README.md +++ b/litellm/proxy/client/README.md @@ -342,7 +342,7 @@ The CLI provides these authentication commands: - **`lite login --pkce`** - Sign in through the system browser with OAuth authorization code + PKCE; the key renews itself with a refresh token - **`lite logout`** - Clear stored authentication token (and revoke a `--pkce` refresh token on the proxy) - **`lite whoami`** - Show current authentication status -- **`lite auth print-token`** - Print the cached token (used as Claude Code's `apiKeyHelper`); renews a `--pkce` key first and fails once a classic token has expired +- **`lite auth print-token`** - Print the cached token (used as Claude Code's `apiKeyHelper`); renews a `--pkce` key first, and fails once a classic token is past 24 hours and the proxy rejects it ### Authentication Flow Steps @@ -382,7 +382,7 @@ Keychain storage needs the `keyring` package, which ships with `pip install 'lit `lite logout` clears both stores. If the keychain is locked at that moment it says so, and re-running it once the keychain is unlocked finishes the job. -The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while fresh and fails once it expires -- there is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. A credential from `lite login --pkce` is the exception: it carries a refresh token, so the CLI renews the key shortly before it expires and `lite logout` revokes the refresh token on the proxy (see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`. +The stored credential is a short-lived, per-session agent token, not a managed virtual key. It is scoped to the user and team you logged in as and inherits their models and budgets; spend is tracked against the shared team and user budgets rather than a separate per-session cap, so multiple logins or several concurrent agents all draw down the same allowance. It is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); re-run `lite login` to refresh it and pick up your latest team and user settings. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it, and once it is past 24 hours (or the client's `LITELLM_CLI_JWT_EXPIRATION_HOURS`) asks the proxy first and fails if the proxy rejects it. There is no silent renewal. It is accepted on a default deployment without `EXPERIMENTAL_UI_LOGIN`, does not appear in the Keys UI, and cannot be rotated or revoked mid-session. A credential from `lite login --pkce` is the exception: it carries a refresh token, so the CLI renews the key shortly before it expires and `lite logout` revokes the refresh token on the proxy (see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. For a long-lived, rotatable, Keys-UI-visible credential, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY`. ### Usage diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 4038ff6f582..819908f9e70 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -506,7 +506,7 @@ To pin the model, pass the agent's own model flag (for example `lite claude --mo The token minted by `lite login` is a short-lived, per-session agent credential, not a managed virtual key. It is scoped to the user and team you authenticated as, inherits that user's and team's models and budgets, and is enforced on the proxy exactly like a virtual key on the same team (guardrails, routing, logging, spend). Spend is tracked against the shared team and user budgets, so running several agents (or logging in more than once) does not hand each session its own separate budget; they all draw down the same team/user allowance. There is no separate per-session cap, so sustained agent use is not capped at a small chat-session limit. -The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the daily re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead. +The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it without asking the proxy for its first 24 hours, or `LITELLM_CLI_JWT_EXPIRATION_HOURS` if you set it on the client. After that, `print-token` and `lite whoami` ask the proxy whether it still accepts the token, and `print-token` fails once the proxy rejects it. There is no silent renewal, so a long-running session needs a fresh `lite login` each time the proxy's token lifetime runs out. If the proxy's tokens last less than 24 hours, set the same `LITELLM_CLI_JWT_EXPIRATION_HOURS` on the client so `print-token` stops printing the token once it runs out. `lite up` still starts a new login once the token is older than that same client window. It copies the token into Claude Code's settings as a static key that it never renews, so a fresh login gives Claude Code the proxy's full token lifetime instead of whatever an old token has left. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead. ### Route Every Claude Code Session Through the Proxy diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index b5302d1f561..e38f54adc4f 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -13,7 +13,6 @@ from rich.console import Console from rich.table import Table from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never -from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, SYSTEM_KEYRING, @@ -143,6 +142,9 @@ _TeamMapping: Final = TypeVar("_TeamMapping", bound=Mapping[str, object]) _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +_PROXY_KEY_CHECK_TIMEOUT_SECONDS: Final = 3 +_PROXY_REJECTED_KEY_STATUSES: Final = frozenset({401, 403}) + KEYRING_INSTALL_HINT: Final = "pip install 'litellm[cli]'" KEYRING_ENABLE_HINT: Final = "keyring --enable (or unset PYTHON_KEYRING_BACKEND)" @@ -1033,10 +1035,12 @@ def print_token(ctx: click.Context): Designed to be used as Claude Code's `apiKeyHelper` (https://docs.claude.com/en/docs/claude-code/settings): stdout must - contain only the token, so all diagnostics go to stderr. The token - expires after `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h); a - `lite login --pkce` token renews itself here first, and once a token - has expired for good, run the same `lite login` command again. + contain only the token, so all diagnostics go to stderr. A token younger + than `LITELLM_CLI_JWT_EXPIRATION_HOURS` (default 24h) is printed without + asking the proxy. An older one is checked against the proxy and refused + only if the proxy rejects it. A `lite login --pkce` + token renews itself here first, and once a token has expired for good, + run the same `lite login` command again. """ vault: Final = context_secret_vault(ctx) token_data: Final = load_token(vault=vault) @@ -1055,7 +1059,7 @@ def print_token(ctx: click.Context): sys.exit(1) renews: Final = "refresh_token" in token_data - if not is_cli_token_fresh(token_data) and not renews: + if not renews and token_expired(token_data): click.echo("Token expired. Run 'lite login' again.", err=True) sys.exit(1) @@ -1111,8 +1115,32 @@ def whoami(ctx: click.Context): expires_at: Final = token_data.get("expires_at") if isinstance(expires_at, (int, float)): click.echo(_key_expiry_line(expires_at, renews="refresh_token" in token_data)) - elif age_hours > CLI_JWT_EXPIRATION_HOURS: - click.echo(f"Warning: Token is more than {CLI_JWT_EXPIRATION_HOURS} hours old and may have expired.") + elif token_expired(token_data): + click.echo("Token expired. Run 'lite login' again.") + + +def token_expired(token_data: Mapping[str, object]) -> bool: + if is_cli_token_fresh(token_data): + return False + if isinstance(token_data.get("expires_at"), (int, float)): + return True + return _proxy_rejects_key(token_data) + + +def _proxy_rejects_key(token_data: Mapping[str, object]) -> bool: + key: Final = token_data.get("key") + base_url: Final = token_data.get("base_url") + if not isinstance(key, str) or not isinstance(base_url, str) or not base_url: + return True + try: + response: Final = requests.get( + f"{base_url}/v1/models", + headers={"Authorization": f"Bearer {key}"}, + timeout=_PROXY_KEY_CHECK_TIMEOUT_SECONDS, + ) + except requests.RequestException: + return False + return response.status_code in _PROXY_REJECTED_KEY_STATUSES def _key_expiry_line(expires_at: float, renews: bool) -> str: diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 1361a3e6364..859322f061d 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -3,9 +3,12 @@ from collections.abc import Mapping from typing import Any, Final import requests +from pydantic import ConfigDict, TypeAdapter from .exceptions import NotFoundError, UnauthorizedError +_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + class ModelsManagementClient: def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): @@ -109,7 +112,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -145,7 +148,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) @@ -287,7 +290,7 @@ class ModelsManagementClient: try: response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() - return response.json() + return _JSON_OBJECT.validate_python(response.json()) except requests.exceptions.HTTPError as e: if e.response.status_code == 401: raise UnauthorizedError(e) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 230d74bc731..418e4002289 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,6 +4,7 @@ import json import logging import math from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime from functools import lru_cache from itertools import chain @@ -137,10 +138,10 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) -from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect +from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict -from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model from litellm.types.extraction import ExtractionResponse from litellm.types.guardrails import GuardrailEventHooks @@ -266,6 +267,8 @@ from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( AnthropicStreamModelRestamper, ) from litellm.proxy.litellm_pre_call_utils import ( + ROUTER_SCHEDULED_ROUTE_TYPES, + SCHEDULER_PRIORITY_HEADER, LiteLLMProxyRequestSetup, add_litellm_data_to_request, get_metadata_variable_name, @@ -2007,6 +2010,49 @@ def _content_text(fields: Mapping[str, object]) -> str: return str(fields.get("content", "")) +async def _no_stream_failure_to_release() -> None: + return None + + +@dataclass(slots=True) +class _ClientStreamFailureReport: + logging_obj: LiteLLMLoggingObj + report: Callable[[Exception], Awaitable[None]] + held: Exception | None = None + + @classmethod + def hold_while_streaming(cls, logging_obj: LiteLLMLoggingObj | None) -> Callable[[], Awaitable[None]]: + if logging_obj is None or logging_obj.on_detached_stream_failure is None: + return _no_stream_failure_to_release + return cls(logging_obj, logging_obj.on_detached_stream_failure)._start_holding() + + def _start_holding(self) -> Callable[[], Awaitable[None]]: + self.logging_obj.on_detached_stream_failure = self.hold + return self.release + + async def hold(self, failure: Exception) -> None: + if self.held is None: + self.held = failure + + async def release(self) -> None: + self.logging_obj.on_detached_stream_failure = self.report + if self.held is None: + return + await asyncio.shield( + spawn_logging_task( + asyncio.get_running_loop(), self._report_held(self.held), task_name="stream failure report" + ) + ) + + async def _report_held(self, failure: Exception) -> None: + try: + await self.report(failure) + except Exception as hook_failure: # noqa: BLE001 # a failing failure hook must not skip the stream cleanup + verbose_proxy_logger.warning( + "Stream failure hook raised: %s(%s)", type(hook_failure).__name__, sanitize_for_log(hook_failure) + ) + + class ProxyBaseLLMRequestProcessing: def __init__(self, data: dict, served_model: str | None = None): self.data = data @@ -2435,14 +2481,20 @@ class ProxyBaseLLMRequestProcessing: self.data["model"] = alias_target aliased_model: Final = self.data.get("model") + if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str): + self.data.update( + llm_router.priority_header_update( + model=aliased_model, kwargs=self.data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER) + ) + ) litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) self.data["litellm_call_id"] = litellm_call_id bind_telemetry_call_id(litellm_call_id) - DDSpanTagger.tag_call_id(self.data.get("litellm_call_id")) + DDSpanTagger.tag_call_id(litellm_call_id) DDSpanTagger.tag_request( user_api_key_dict=user_api_key_dict, - requested_model=self.data.get("model"), + requested_model=aliased_model, ) ### AUTO STREAM USAGE TRACKING ### @@ -3508,6 +3560,8 @@ class ProxyBaseLLMRequestProcessing: request_data: Final = self.data async def _on_detached_stream_failure(exc: Exception) -> None: + if not logging_obj.claim_stream_failure_report(): + return try: await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -4458,11 +4512,15 @@ class ProxyBaseLLMRequestProcessing: client_disconnected = False delivered_chunk = False recent_tail = SSE_STREAM_START_TAIL # rebind-ok: rolling window over the yielded bytes + stream_logging_obj: Final[object] = request_data.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # legacy request data is narrowed to Logging below guarded_stream: Final[AsyncGenerator[object, None]] = proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=response, request_data=request_data, ) + release_stream_failure: Final = _ClientStreamFailureReport.hold_while_streaming( + stream_logging_obj if isinstance(stream_logging_obj, LiteLLMLoggingObj) else None + ) try: str_so_far = "" async for chunk in guarded_stream: @@ -4533,10 +4591,18 @@ class ProxyBaseLLMRequestProcessing: verbose_proxy_logger.exception( "litellm.proxy.proxy_server.async_data_generator(): Exception occured - %s", e ) - transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=request_data, + already_reported: Final = ( + isinstance(stream_logging_obj, LiteLLMLoggingObj) + and not stream_logging_obj.claim_stream_failure_report() + ) + transformed_exception: Final = ( + None + if already_reported + else await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=e, + request_data=request_data, + ) ) if transformed_exception is not None: e = transformed_exception @@ -4559,15 +4625,18 @@ class ProxyBaseLLMRequestProcessing: seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail) yield seal + error_frame if seal else error_frame finally: - await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( - request=request, - request_data=request_data, - response=response, - stream_completed=stream_completed, - client_disconnected=client_disconnected, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - ) + try: + await release_stream_failure() + finally: + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=request, + request_data=request_data, + response=response, + stream_completed=stream_completed, + client_disconnected=client_disconnected, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) @staticmethod def async_sse_data_generator( diff --git a/litellm/proxy/db/accounting_ledger.py b/litellm/proxy/db/accounting_ledger.py new file mode 100644 index 00000000000..91f8307283e --- /dev/null +++ b/litellm/proxy/db/accounting_ledger.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import hashlib +import json +from collections.abc import AsyncGenerator, Callable +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from pydantic import AwareDatetime, ConfigDict, JsonValue, TypeAdapter +from typing_extensions import assert_never + +from litellm.accounting import CostEvent, DeliveryTarget, EventKind, InboundUsageEvent, OutboundUsageEvent +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper + +if TYPE_CHECKING: + from prisma import Prisma + + from litellm.proxy.utils import PrismaClient + +AccountingEvent: TypeAlias = InboundUsageEvent | OutboundUsageEvent | CostEvent +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_INSTANT: Final[TypeAdapter[datetime]] = TypeAdapter(AwareDatetime) +_STRINGS: Final = TypeAdapter(tuple[dict[str, str], ...]) +_TIMEOUT: Final = timedelta(seconds=10) + + +class AccountingConflictError(ValueError): + pass + + +_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"] + + +@dataclass(frozen=True, slots=True) +class _AcceptedEvent: + event_id: str + kind: EventKind + source_event_id: str | None + occurred_at: datetime + payload: str + payload_hash: str + targets: tuple[str, ...] + schema_version: int + version: int + + +def _json(value: object) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False) + + +def _prepare( + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...], + envelope: str, + schema_version: int, + version: int, +) -> _AcceptedEvent: + if type(schema_version) is not int or type(version) is not int or min(schema_version, version) < 1: + raise ValueError("Accounting versions must be positive integers") + validated: Final = type(event).model_validate_json(event.model_dump_json()) + payload: Final = _JSON.validate_json(envelope) + if payload.get("event") != validated.model_dump(mode="json"): + raise ValueError("Accounting envelope must contain the matching public event") + destinations: Final = tuple(sorted(frozenset(DeliveryTarget(target).value for target in targets))) + serialized: Final = _json(payload) + identity: Final = _json((schema_version, version, destinations, payload)) + kind: Final = ( + EventKind.COST + if isinstance(validated, CostEvent) + else EventKind.INBOUND_USAGE + if isinstance(validated, InboundUsageEvent) + else EventKind.OUTBOUND_USAGE + ) + return _AcceptedEvent( + event_id=validated.event_id, + kind=kind, + source_event_id=validated.usage_event_id if isinstance(validated, CostEvent) else None, + occurred_at=validated.occurred_at if isinstance(validated, CostEvent) else validated.ended_at, + payload=serialized, + payload_hash=hashlib.sha256(identity.encode()).hexdigest(), + targets=destinations, + schema_version=schema_version, + version=version, + ) + + +_INSERT_EVENT: Final = """ +INSERT INTO "LiteLLM_AccountingEvent" + (event_id, kind, source_event_id, schema_version, payload_hash, occurred_at, payload) +VALUES ($1, $2::"AccountingEventKind", $3, $4::int, $5, $6::timestamp, $7::jsonb) +ON CONFLICT (event_id) DO NOTHING RETURNING event_id +""" +_INSERT_DELIVERIES: Final = """ +INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) +SELECT $1, value::"AccountingDeliveryTarget", $3::int FROM jsonb_array_elements_text($2::jsonb) +""" +_LOCK_SOURCE: Final = """ +SELECT event_id FROM "LiteLLM_AccountingEvent" +WHERE event_id = $1 AND kind = 'outbound_usage' FOR UPDATE +""" +_EXISTING_COST: Final = """ +SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_event_id = $1 +""" +_LOCK_EVENT: Final = """ +SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE +""" +_COMPLETED_FAMILY: Final = """ +root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage') +AND root.recorded_at < $1::timestamp +AND (root.kind = 'inbound_usage' OR EXISTS ( + SELECT 1 FROM "LiteLLM_AccountingEvent" AS cost WHERE cost.source_event_id = root.event_id AND cost.kind = 'cost' +)) +AND NOT EXISTS ( + SELECT 1 FROM "LiteLLM_AccountingEvent" AS member + WHERE (member.event_id = root.event_id OR (member.kind = 'cost' AND member.source_event_id = root.event_id)) + AND (member.recorded_at >= $1::timestamp + OR EXISTS (SELECT 1 FROM "LiteLLM_AccountingDelivery" AS delivery + WHERE delivery.event_id = member.event_id AND delivery.state != 'applied')) +) +""" +_RETENTION_CANDIDATES: Final = ( + 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE ' + + _COMPLETED_FAMILY + + " ORDER BY root.recorded_at, root.event_id LIMIT $2::int FOR UPDATE OF root SKIP LOCKED" +) +_ELIGIBLE_GROUP: Final = ( + 'SELECT root.event_id FROM "LiteLLM_AccountingEvent" AS root WHERE root.event_id = $2 AND ' + _COMPLETED_FAMILY +) +_DELETE_DELIVERIES: Final = """ +DELETE FROM "LiteLLM_AccountingDelivery" WHERE event_id IN ( + SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) +) +""" +_DELETE_EVENTS: Final = """ +DELETE FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 OR (kind = 'cost' AND source_event_id = $1) +""" + + +_Conflict: TypeAlias = Literal["missing_source", "competing_cost", "identity_conflict"] + + +async def _accept(db: Prisma, record: _AcceptedEvent) -> Literal["inserted", "duplicate"] | _Conflict: + if record.source_event_id is not None: + source: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_SOURCE, record.source_event_id)) + if not source: + return "missing_source" + costs: Final = _STRINGS.validate_python(tuple(await db.query_raw(_EXISTING_COST, record.source_event_id))) + if any(cost["event_id"] != record.event_id for cost in costs): + return "competing_cost" + inserted: Final = _STRINGS.validate_python( + await db.query_raw( + _INSERT_EVENT, + record.event_id, + record.kind.value, + record.source_event_id, + record.schema_version, + record.payload_hash, + record.occurred_at.astimezone(timezone.utc).isoformat(), + record.payload, + ) + ) + if not inserted: + existing: Final = _STRINGS.validate_python(tuple(await db.query_raw(_LOCK_EVENT, record.event_id))) + if not existing or existing[0]["payload_hash"] != record.payload_hash: + return "identity_conflict" + return "duplicate" + await db.execute_raw(_INSERT_DELIVERIES, record.event_id, _json(record.targets), record.version) + return "inserted" + + +async def _delete_completed(db: Prisma, event_id: str, cutoff: str) -> int: + if not _STRINGS.validate_python(await db.query_raw(_ELIGIBLE_GROUP, cutoff, event_id)): + return 0 + await db.execute_raw(_DELETE_DELIVERIES, event_id) + return await db.execute_raw(_DELETE_EVENTS, event_id) + + +class AccountingLedger: + def __init__(self, transaction: Callable[[], _Transaction]) -> None: + self._transaction: Final = transaction + + @classmethod + def for_client(cls, client: PrismaClient) -> AccountingLedger: + from prisma.client import TransactionManager + + managers: Final = TypeAdapter(TransactionManager, config=ConfigDict(arbitrary_types_allowed=True)) + + @asynccontextmanager + async def transaction() -> AsyncGenerator[Prisma]: + primary: Final = writer_wrapper(client.db) + manager: Final = managers.validate_python(primary.tx(timeout=_TIMEOUT)) + async with db_span("accounting_ledger", "LiteLLM_AccountingEvent"), manager as db: + await db.execute_raw("SET LOCAL statement_timeout = 5000") + await db.execute_raw("SET LOCAL lock_timeout = 1000") + yield db + + return cls(transaction) + + async def insert_event( + self, + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...], + *, + envelope: str, + schema_version: int = 1, + version: int = 1, + ) -> Literal["inserted", "duplicate"]: + """Accept T5/T6's serialized envelope, containing its public event under `event`. + + The immutable acceptance hash includes the target set and delivery version. + Changing either for an accepted event conflicts. T7 must provide a separate + add-target or new-version workflow that preserves receipts and prevents + duplicate contributions; redelivery cannot silently create or reset work. + + The Prisma adapter opens one interactive transaction per event, including + two SET LOCAL statements and a source-row lock for costs. I/O is async, + but awaiting this method inline adds commit latency per event, potentially + several times per request. T5 owns bounded asynchronous capture and an + explicit batching strategy; looping over this method does not batch commits. + Success acknowledges commit, independently of HTTP response acknowledgement. + """ + record: Final = _prepare(event, targets, envelope, schema_version, version) + async with self._transaction() as db: + result: Final = await _accept(db, record) + match result: + case "inserted" | "duplicate": + return result + case "missing_source": + raise AccountingConflictError("Cost source must be a durably accepted outbound usage event") + case "competing_cost": + raise AccountingConflictError("Usage event already has an authoritative cost") + case "identity_conflict": + raise AccountingConflictError("Event identity conflicts with its accepted content or delivery set") + case _: + assert_never(result) + + async def purge_completed( + self, + *, + redelivery_before: datetime | None = None, + recovery_before: datetime | None = None, + limit: int = 100, + ) -> int: + """Delete at most `limit` completed families once both retention owners permit it. + + T8 must supply a coherent recovery watermark including outstanding holds. + Omitting either watermark disables deletion. Each usage row and its cost + remain together so a retained charge never loses its source or vice versa. + Applied receipts are deleted with the family, ending its retained dedup + protection. The caller must guarantee events recorded before + redelivery_before cannot be redelivered through supported paths. + T8's recovery_before must also protect checkpoints, + outstanding holds and unfinished settlements; age alone is insufficient. + """ + if type(limit) is not int or not 1 <= limit <= 1000: + raise ValueError("Accounting retention batch size must be between 1 and 1000") + if redelivery_before is None or recovery_before is None: + return 0 + cutoff: Final = min( + _INSTANT.validate_python(redelivery_before).astimezone(timezone.utc), + _INSTANT.validate_python(recovery_before).astimezone(timezone.utc), + ).isoformat() + async with self._transaction() as db: + roots: Final = _STRINGS.validate_python(tuple(await db.query_raw(_RETENTION_CANDIDATES, cutoff, limit))) + deleted: Final = tuple([await _delete_completed(db, root["event_id"], cutoff) for root in roots]) + return sum(deleted) diff --git a/litellm/proxy/db/budget_window_spend_writer.py b/litellm/proxy/db/budget_window_spend_writer.py index 82c95cfe998..99b49331932 100644 --- a/litellm/proxy/db/budget_window_spend_writer.py +++ b/litellm/proxy/db/budget_window_spend_writer.py @@ -64,6 +64,16 @@ _ROLL_WINDOW_SPEND_SQL: Final = ( "AND window_start < ($4::timestamptz AT TIME ZONE 'UTC')" ) +_OPEN_WINDOW_SPEND_SQL: Final = ( + 'INSERT INTO "LiteLLM_BudgetWindowSpend" ' + "(entity_type, entity_id, window_duration, window_start, spend, created_at, updated_at) " + "VALUES ($1, $2, $3, ($4::timestamptz AT TIME ZONE 'UTC'), 0, " + "($4::timestamptz AT TIME ZONE 'UTC'), ($4::timestamptz AT TIME ZONE 'UTC')) " + "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE SET " + "window_start = EXCLUDED.window_start, spend = 0, updated_at = EXCLUDED.updated_at " + "WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')" +) + _SEED_FROM_SPEND_LOGS_KEY_SQL: Final = ( "SELECT COALESCE(SUM(spend), 0.0) AS total, " "COALESCE(SUM(spend) FILTER (WHERE \"startTime\" < ($3::timestamptz AT TIME ZONE 'UTC')), 0.0) AS before_batch " @@ -389,3 +399,29 @@ async def roll_window_spend_row( to_naive_utc(new_window_start), to_naive_utc(datetime.now(timezone.utc)), ) + + +async def open_window_spend_row( + prisma_client: "PrismaClient", + entity_type: str, + entity_id: str, + window_duration: str, + window_start: datetime, + replaces_window_start: datetime | None, +) -> None: + """Start a window at `window_start` unless another writer already started one. + + The row is created when it does not exist, and moved onto the new window only + while it still holds `replaces_window_start`, the ended window the caller saw. + A writer that lost the race leaves the winner's window in place, so reading + the row back gives every writer the same window start. + """ + async with db_span("open_window_spend_row", "LiteLLM_BudgetWindowSpend"): + await prisma_client.db.execute_raw( + _OPEN_WINDOW_SPEND_SQL, + entity_type, + entity_id, + window_duration, + to_naive_utc(window_start), + None if replaces_window_start is None else to_naive_utc(replaces_window_start), + ) diff --git a/litellm/proxy/google_endpoints/agents_endpoints.py b/litellm/proxy/google_endpoints/agents_endpoints.py index 00b23d846c2..7e3c2ddafe0 100644 --- a/litellm/proxy/google_endpoints/agents_endpoints.py +++ b/litellm/proxy/google_endpoints/agents_endpoints.py @@ -160,9 +160,10 @@ async def create_gemini_agent( -H "Authorization: Bearer sk-..." \\ -H "Content-Type: application/json" \\ -d '{ - "name": "my-custom-slides-agent", - "base_agent": "waverunner", - "instructions": "You are a helpful assistant that creates slides.", + "id": "my-custom-slides-agent", + "base_agent": "antigravity-preview-09-2026", + "system_instruction": "You are a helpful assistant that creates slides.", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, "base_environment": { "type": "remote", "sources": [ diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index b4830b7c771..9bea22acd71 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -305,6 +305,11 @@ def iter_message_text(data: Mapping[str, object]) -> Iterator[str]: yield from _iter_text_parts_in_content(message.get("content")) +def _is_tool_message(message: object) -> bool: + fields: Final = _text_object(message) + return fields is not None and str(fields.get("role") or "").lower() == "tool" + + def _call_item(item: object) -> tuple[str, object] | None: """Return ``(field, value)`` for a Responses tool-call input item, else None.""" fields: Final = _text_object(item) @@ -316,11 +321,14 @@ def walk_user_text( data: dict[str, object], visit: Callable[[str], str], visit_tool: Callable[[JsonValue], JsonValue] | None = None, + *, + skip_tool_message: bool = False, ) -> int: """Rewrite every text fragment in place via ``visit``. Mutates ``data["messages"]`` and ``data["input"]`` (including tool payloads), ``data["system"]`` and - ``data["instructions"]``. Returns the number + ``data["instructions"]``. ``skip_tool_message`` leaves ``role: tool`` messages and ``tool_result`` + blocks as sent. Returns the number of fragments visited so callers can short-circuit when nothing was inspected. """ @@ -353,6 +361,8 @@ def walk_user_text( def _rewrite_part(part: Mapping[str, object], native: bool) -> Mapping[str, object]: nonlocal visited + if skip_tool_message and part.get("type") == "tool_result": + return part text: Final = _part_text(part) rewritten: Final = {**part, **({"text": visit(text)} if text is not None else {})} if text is not None: @@ -384,18 +394,17 @@ def walk_user_text( return content messages: Final = data.get("messages") - if isinstance(messages, list): - for message in messages: - if not isinstance(message, dict): - continue - if "content" in message: - message["content"] = _rewrite_content(message["content"]) - calls = message.get("tool_calls") - if isinstance(calls, list): - for call in calls: - function = call.get("function") if isinstance(call, dict) else None - if isinstance(function, dict) and "arguments" in function: - function["arguments"] = _rewrite_arguments(function["arguments"]) + for message in _object_list(messages) or (): + if (skip_tool_message and _is_tool_message(message)) or not isinstance(message, dict): + continue + if "content" in message: + message["content"] = _rewrite_content(message["content"]) + calls = message.get("tool_calls") + if isinstance(calls, list): + for call in calls: + function = call.get("function") if isinstance(call, dict) else None + if isinstance(function, dict) and "arguments" in function: + function["arguments"] = _rewrite_arguments(function["arguments"]) if "system" in data: data["system"] = _rewrite_content(data["system"]) # rebind-ok: preserve the native hook's in-place contract diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 1267fb15db8..3071bcbb1dc 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -14,7 +14,7 @@ from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket -from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX +from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 98be35079c3..16b3c45f6f8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -169,7 +169,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): from litellm.types.llms.openai import OpenAIModerationResponse - return OpenAIModerationResponse(**response.json()) + return OpenAIModerationResponse.model_validate(response.json()) def _check_moderation_result(self, moderation_response: "OpenAIModerationResponse") -> None: """ diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 0d2f8dc9a0c..be32405154c 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -27,7 +27,7 @@ from litellm.constants import ( HEALTH_CHECK_TIMEOUT_SECONDS, ) from litellm.litellm_core_utils.health_check_helpers import native_health_check_mode -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( StrategyRouterDependency, classify_strategy_router_model, strategy_router_dependencies, diff --git a/litellm/proxy/hooks/autorouter_baseline_cache.py b/litellm/proxy/hooks/autorouter_baseline_cache.py index c4232609124..1147ce7ec12 100644 --- a/litellm/proxy/hooks/autorouter_baseline_cache.py +++ b/litellm/proxy/hooks/autorouter_baseline_cache.py @@ -37,8 +37,8 @@ from litellm.proxy.spend_tracking.savings import ( _effective_model_info, # pyright: ignore[reportPrivateUsage] # existing deployment-price owner _proxy_llm_router, # pyright: ignore[reportPrivateUsage] # existing optional proxy-router owner ) +from litellm.router_strategy.auto_router.baseline_request import baseline_request from litellm.router_strategy.complexity_router.context_compaction import compaction_applied -from litellm.router_utils.baseline_request import baseline_request from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import BaselineRouteStamp from litellm.types.utils import CallTypes, ModelInfo, Usage diff --git a/litellm/proxy/hooks/model_budget_shared_windows.py b/litellm/proxy/hooks/model_budget_shared_windows.py new file mode 100644 index 00000000000..880dd0eb218 --- /dev/null +++ b/litellm/proxy/hooks/model_budget_shared_windows.py @@ -0,0 +1,176 @@ +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final, Protocol, TypeAlias + +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.proxy.db.budget_window_spend_writer import open_window_spend_row +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + WindowSpendTransaction, + build_window_spend_transaction, +) +from litellm.repositories.table_repositories import BudgetWindowSpendRepository +from litellm.router_strategy.budget_limiter import budget_window_seconds + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +MODEL_BUDGET_WINDOW_ENTITY_TYPE: Final = "model_max_budget" +_READ_TTL_SECONDS: Final = 5 + + +@dataclass(frozen=True, slots=True) +class SharedWindow: + started_at: float + spend: float + + def is_open(self, budget_duration: str, now: float) -> bool: + return now < self.started_at + budget_window_seconds(self.started_at, budget_duration) + + +class LocalTTLCache(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: object, *, ttl: float) -> None: ... + + def delete_cache(self, key: str) -> None: ... + + +class SharedWindows(Protocol): + async def read(self, spend_key: str, budget_duration: str) -> SharedWindow | None: ... + + async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None: ... + + async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None: ... + + +PrismaClientSource: TypeAlias = Callable[[], "PrismaClient | None"] +WindowSpendSink: TypeAlias = Callable[[WindowSpendTransaction], Awaitable[None]] + + +def _proxy_prisma_client() -> "PrismaClient | None": + try: + from litellm.proxy.proxy_server import prisma_client + except ImportError: + return None + return prisma_client + + +async def _queue_on_proxy_spend_writer(transaction: WindowSpendTransaction) -> None: + from litellm.proxy.proxy_server import proxy_logging_obj + + await proxy_logging_obj.db_spend_update_writer.window_spend_update_queue.add_update(transaction) + + +@dataclass(frozen=True, slots=True) +class _Reading: + window: SharedWindow | None + + +class DatabaseSharedWindows: + def __init__( + self, + prisma_client_source: PrismaClientSource = _proxy_prisma_client, + sink: WindowSpendSink = _queue_on_proxy_spend_writer, + read_ttl_seconds: int = _READ_TTL_SECONDS, + ) -> None: + self._prisma_client_source = prisma_client_source + self._sink = sink + self._read_ttl_seconds = read_ttl_seconds + self._readings: LocalTTLCache = InMemoryCache(max_size_in_memory=10_000, default_ttl=read_ttl_seconds) + + async def read(self, spend_key: str, budget_duration: str) -> SharedWindow | None: + cached: Final = self._readings.get_cache(key=spend_key) + if isinstance(cached, _Reading): + return cached.window + window: Final = await self._read_row(spend_key=spend_key, budget_duration=budget_duration) + self._readings.set_cache(key=spend_key, value=_Reading(window=window), ttl=self._read_ttl_seconds) + return window + + async def open(self, spend_key: str, budget_duration: str, now: float) -> float | None: + """When the window every replica shares started, opening one at `now` if none is open. + + Without a database every window is this replica's own, so it starts at `now`. None means the database did + not answer, so this replica cannot know which window the others are in yet. + """ + prisma_client: Final = self._prisma_client_source() + if prisma_client is None: + return now + try: + seen: Final = await _fetch_row( + prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=True + ) + if seen is not None and seen.is_open(budget_duration=budget_duration, now=now): + return seen.started_at + await open_window_spend_row( + prisma_client=prisma_client, + entity_type=MODEL_BUDGET_WINDOW_ENTITY_TYPE, + entity_id=spend_key, + window_duration=budget_duration, + window_start=datetime.fromtimestamp(now, tz=timezone.utc), + replaces_window_start=None + if seen is None + else datetime.fromtimestamp(seen.started_at, tz=timezone.utc), + ) + opened: Final = await _fetch_row( + prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=True + ) + except Exception: # noqa: BLE001 # an unreachable database leaves this replica on its own window + verbose_proxy_logger.exception("Unable to open the shared model budget window") + return None + if opened is None or not opened.is_open(budget_duration=budget_duration, now=now): + return None + return opened.started_at + + async def add(self, spend_key: str, budget_duration: str, started_at: float, spend: float) -> None: + if self._prisma_client_source() is None: + return + try: + await self._sink( + build_window_spend_transaction( + entity_type=MODEL_BUDGET_WINDOW_ENTITY_TYPE, + entity_id=spend_key, + window_duration=budget_duration, + window_start=datetime.fromtimestamp(started_at, tz=timezone.utc), + spend=spend, + ) + ) + except Exception: # noqa: BLE001 # spend tracking must never fail the success callback + verbose_proxy_logger.debug("Unable to queue shared model budget spend", exc_info=True) + + async def _read_row(self, spend_key: str, budget_duration: str) -> SharedWindow | None: + prisma_client: Final = self._prisma_client_source() + if prisma_client is None: + return None + try: + return await _fetch_row( + prisma_client=prisma_client, spend_key=spend_key, budget_duration=budget_duration, use_writer=False + ) + except Exception: # noqa: BLE001 # an unreadable row leaves enforcement on this replica's own counter + verbose_proxy_logger.exception("Unable to read shared model budget spend") + return None + + +async def _fetch_row( + prisma_client: "PrismaClient", spend_key: str, budget_duration: str, use_writer: bool +) -> SharedWindow | None: + row: Final = await BudgetWindowSpendRepository(prisma_client, use_writer=use_writer).table.find_unique( + where={ + "entity_type_entity_id_window_duration": { + "entity_type": MODEL_BUDGET_WINDOW_ENTITY_TYPE, + "entity_id": spend_key, + "window_duration": budget_duration, + } + } + ) + if row is None: + return None + window_start: Final = row.window_start + started_at: Final = ( + window_start if window_start.tzinfo is not None else window_start.replace(tzinfo=timezone.utc) + ).timestamp() + return SharedWindow(started_at=started_at, spend=float(row.spend or 0.0)) + + +PROXY_SHARED_WINDOWS: Final = DatabaseSharedWindows() diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index bc9cee1cb1d..1c4a61cfad8 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,25 +1,41 @@ import asyncio +import functools +import itertools import json import time -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType -from typing import Final, cast # noqa: TID251 # untyped callback payload requires explicit narrowing +from typing import ( + TYPE_CHECKING, + Final, + TypeAlias, + cast, # noqa: TID251 # untyped callback payload requires explicit narrowing +) import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span +from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup -from litellm.llms.bedrock.common_utils import get_bedrock_base_model +from litellm.llms.bedrock.common_utils import get_bedrock_base_model, get_bedrock_cross_region_inference_regions from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth +from litellm.proxy.hooks.model_budget_shared_windows import ( + PROXY_SHARED_WINDOWS, + LocalTTLCache, + SharedWindow, + SharedWindows, +) from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds -from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.types.llms.openai import AllMessageValues -from litellm.types.utils import BudgetConfig, StandardLoggingPayload +from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload + +if TYPE_CHECKING: + from litellm.router import Router VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend" END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend" @@ -39,10 +55,18 @@ _SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType( _LEGACY_REQUEST_MODEL_SCOPES: Final = frozenset({Litellm_EntityType.KEY, Litellm_EntityType.END_USER}) +_PER_REPLICA_SCOPES: Final = frozenset({Litellm_EntityType.END_USER}) + _PROCESS_STARTED_AT: Final = time.monotonic() _LATEST_RESET_AT: Final = datetime(9999, 12, 31, tzinfo=timezone.utc).timestamp() +_GLOB_WILDCARD: Final = "*" + +ModelAccessGroups: TypeAlias = Callable[[str, str | None], Collection[str]] + +_PROVIDER_PREFIXES: Final = tuple(f"{provider.value}/" for provider in LlmProviders) + _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType( { Litellm_EntityType.KEY: "virtual_key_budget_start_time", @@ -68,6 +92,15 @@ class ResolvedModelBudget: budget_config: BudgetConfig +@dataclass(frozen=True, slots=True) +class RoutedModels: + served_names: frozenset[str] + group_members: Mapping[str, frozenset[str]] + + def members(self, access_group: str) -> frozenset[str]: + return self.group_members.get(access_group, frozenset()) + + def team_member_budget_entity_id(user_id: str, team_id: str) -> str: return f"{len(user_id)}:{user_id}:{team_id}" @@ -131,43 +164,403 @@ def model_budget_start_time_cache_key( return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}" -def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None: - """Find the `model_max_budget` entry that governs `model`, or None.""" - for candidate in _budget_model_candidates(model): - raw_budget_config = model_max_budget.get(candidate) - if raw_budget_config is None: - continue - if (budget_config := _usable_budget_config(raw_budget_config)) is None: - # An entry that will not validate cannot be keyed, so it cannot be - # enforced or incremented. Skip to the next candidate rather than - # raising: raising would abort every other scope's increment and turn - # a config typo into a 500, and stopping here would let one malformed - # specific entry disable a perfectly good bare-family budget beside - # it. The candidate chain already falls through an ABSENT entry, and - # an unparseable one is indistinguishable from absent to enforcement. - # `validate_model_max_budget` rejects these on the write path, so - # reaching here means config.yaml or a direct DB edit. - verbose_proxy_logger.warning( - "Ignoring unusable model_max_budget entry for %s; it cannot be enforced or tracked", - candidate, - ) - continue - return ResolvedModelBudget(budget_model=candidate, budget_config=budget_config) - return None +def model_access_groups_of(router: "Router", model: str, team_id: str | None = None) -> frozenset[str]: + """Access groups serving `model`, or the team's own model of that name, leaving out any named after a model the + router serves. + + An entry with that name already budgets that model, so reading it as the group too would quietly widen it. + """ + return frozenset( + group + for group in router.get_model_access_groups(model_name=model, team_id=team_id) + if group not in router.model_names + ) -def team_model_budget_applies(model: str, key_model_max_budget: Mapping[str, object] | None) -> bool: - """A key entry that spend-gates `model` overrides the team cap: it is then gated on and billed to the key alone.""" +def routed_models_of(router: "Router", team_id: str | None = None) -> RoutedModels: + serves_every_name: Final = ( + router.default_deployment is not None or router.router_general_settings.pass_through_all_models + ) + served_names: Final = frozenset(router.model_names) | ( + frozenset({_GLOB_WILDCARD}) if serves_every_name else frozenset() + ) + team_members: Final = _team_public_names_by_group(router=router, team_id=team_id) + return RoutedModels( + served_names=served_names, + group_members=MappingProxyType( + { + group: frozenset(members) | team_members.get(group, frozenset()) + for group, members in router.get_model_access_groups().items() + if group not in served_names + } + ), + ) + + +def _team_public_names_by_group(router: "Router", team_id: str | None) -> Mapping[str, frozenset[str]]: + if team_id is None: + return MappingProxyType({}) + groups_by_name: Final = MappingProxyType( + { + name: model_access_groups_of(router=router, model=name, team_id=team_id) + for name in router.team_public_model_names_of(team_id) + } + ) + return MappingProxyType( + { + group: frozenset(name for name, groups in groups_by_name.items() if group in groups) + for group in frozenset[str]().union(*groups_by_name.values()) + } + ) + + +def proxy_model_access_groups(model: str, team_id: str | None = None) -> frozenset[str]: + router: Final = _proxy_router() + return frozenset() if router is None else model_access_groups_of(router=router, model=model, team_id=team_id) + + +def proxy_routed_models(team_id: str | None = None) -> RoutedModels: + router: Final = _proxy_router() + if router is None: + return RoutedModels(served_names=frozenset(), group_members=MappingProxyType({})) + return routed_models_of(router=router, team_id=team_id) + + +def _proxy_router() -> "Router | None": + try: + from litellm.proxy.proxy_server import llm_router + except ImportError: + return None + return llm_router + + +def resolve_model_budget( + model: str, + model_max_budget: Mapping[str, object], + model_access_groups: ModelAccessGroups = proxy_model_access_groups, + team_id: str | None = None, +) -> ResolvedModelBudget | None: + """Find the `model_max_budget` entry that governs `model`, or None. + + An entry named after the model wins, then one named after a model access + group serving it, then the most specific glob (`claude-opus*`). Every model + an access group or glob entry covers draws from that entry's one counter. + """ + resolved: Final = ( + _usable_entry(entry=entry, raw_budget_config=model_max_budget[entry]) + for entry in _governing_entry_names( + model=model, model_max_budget=model_max_budget, model_access_groups=model_access_groups, team_id=team_id + ) + ) + return next((entry for entry in resolved if entry is not None), None) + + +def _governing_entry_names( + model: str, + model_max_budget: Mapping[str, object], + model_access_groups: ModelAccessGroups, + team_id: str | None, +) -> Iterator[str]: + candidates: Final = _budget_model_candidates(model) + yield from (candidate for candidate in candidates if candidate in model_max_budget) + group_entries: Final = tuple( + sorted(entry for entry in model_max_budget if _GLOB_WILDCARD not in entry and entry not in candidates) + ) + if group_entries: + groups: Final = model_access_groups(model, team_id) + yield from (entry for entry in group_entries if entry in groups) + yield from sorted( + (entry for entry in model_max_budget if _GLOB_WILDCARD in entry and _glob_matches(entry, candidates)), + key=_glob_precedence, + reverse=True, + ) + + +def _glob_matches(pattern: str, names: Iterable[str]) -> bool: + return any(_glob_matches_name(pattern=pattern, name=name) for name in names) + + +def _glob_matches_name(pattern: str, name: str) -> bool: + """Whole-name match where `*` is the only wildcard, in linear time whatever the pattern holds. + + Taking each literal piece between two wildcards at its first occurrence is + always safe, since a later occurrence only leaves less room for the rest, so + the match never needs to backtrack the way a `.*` regex does. + """ + if _GLOB_WILDCARD not in pattern: + return pattern == name + head, *middle, tail = pattern.split(_GLOB_WILDCARD) + middle_end: Final = len(name) - len(tail) + if middle_end < len(head) or not name.startswith(head) or not name.endswith(tail): + return False + return functools.reduce(_after_piece(name=name, end=middle_end), middle, len(head)) >= 0 + + +def _after_piece(name: str, end: int) -> Callable[[int, str], int]: + def step(position: int, piece: str) -> int: + found: Final = name.find(piece, position, end) if position >= 0 else -1 + return found + len(piece) if found >= 0 else -1 + + return step + + +def _glob_precedence(pattern: str) -> tuple[int, int, str]: + """More literal characters first, then fewer wildcards, so padding a pattern with `*` never makes it win.""" + wildcards: Final = pattern.count(_GLOB_WILDCARD) + return len(pattern) - wildcards, -wildcards, pattern + + +@dataclass(frozen=True, slots=True) +class _RequestForm: + """Request names that an entry governs. + + Every such name matches all of ``name_globs`` and also resolves through + ``candidate_globs``. Its last ``/`` segment is among them unless + ``last_segment_open``. + """ + + name_globs: tuple[str, ...] + candidate_globs: tuple[str, ...] = () + last_segment_open: bool = False + + +def entries_outranked_by(entry: str, model_max_budget: Collection[str], routed_models: RoutedModels) -> Iterator[str]: + """Entries that `entry` can take precedence over for some model a request can name. + + A model or group entry is resolved before any pattern, and a pattern before + every less specific one, so `entry` outranks a pattern that some request it + governs also resolves to. A group entry with a wildcard route among its + members outranks every pattern. + """ + patterns: Final = tuple(other for other in model_max_budget if _GLOB_WILDCARD in other) + if _GLOB_WILDCARD in entry: + entry_precedence: Final = _glob_precedence(entry) + pattern_forms: Final = tuple(_pattern_request_forms(pattern=entry, routed_models=routed_models)) + yield from ( + other + for other in patterns + if _glob_precedence(other) < entry_precedence and _resolves_any(pattern=other, forms=pattern_forms) + ) + return + entry_members: Final = routed_models.members(entry) + covers_wildcard_route: Final = any(_GLOB_WILDCARD in member for member in entry_members) + entry_forms: Final = tuple( + _entry_request_forms(entry=entry, entry_members=entry_members, routed_models=routed_models) + ) + yield from (other for other in patterns if covers_wildcard_route or _resolves_any(pattern=other, forms=entry_forms)) + yield from ( + other + for other in model_max_budget + if _GLOB_WILDCARD not in other + and _ranks_ahead_of_group( + entry=entry, entry_members=entry_members, group=other, group_members=routed_models.members(other) + ) + ) + + +def _resolves_any(pattern: str, forms: Iterable[_RequestForm]) -> bool: + return any(_resolves(pattern=pattern, form=form) for form in forms) + + +def _resolves(pattern: str, form: _RequestForm) -> bool: + return ( + _globs_share_a_name(pattern, *form.name_globs) + or any(_globs_share_a_name(pattern, candidate) for candidate in form.candidate_globs) + or (form.last_segment_open and "/" not in pattern and _globs_share_a_name(f"*/{pattern}", *form.name_globs)) + ) + + +def _globs_share_a_name(*globs: str) -> bool: + """Whether one name matches every glob. + + Globs that each hold a `*` share a name exactly when their literal heads + are prefixes of one another and so are their literal tails: the longest + head, every middle piece and the longest tail, in that order, matches all + of them. + """ + literal: Final = next((glob for glob in globs if _GLOB_WILDCARD not in glob), None) + if literal is not None: + return all(_glob_matches_name(pattern=glob, name=literal) for glob in globs) + heads: Final = sorted((glob.partition(_GLOB_WILDCARD)[0] for glob in globs), key=len) + tails: Final = sorted((glob.rpartition(_GLOB_WILDCARD)[2] for glob in globs), key=len) + return all(longer.startswith(shorter) for shorter, longer in itertools.pairwise(heads)) and all( + longer.endswith(shorter) for shorter, longer in itertools.pairwise(tails) + ) + + +def _served_form(name: str) -> _RequestForm: + return _RequestForm(name_globs=(name,), candidate_globs=_budget_model_candidates(name)) + + +def _wildcard_routes(routed_models: RoutedModels) -> tuple[str, ...]: + return tuple(sorted(name for name in routed_models.served_names if _GLOB_WILDCARD in name)) + + +def _entry_request_forms( + entry: str, entry_members: frozenset[str], routed_models: RoutedModels +) -> Iterator[_RequestForm]: + yield from ( + _served_form(name) + for name in _names_governed_by(entry=entry, entry_members=entry_members, routed_models=routed_models) + ) + routes: Final = _wildcard_routes(routed_models) + yield from ( + _RequestForm(name_globs=(route, f"*/{entry}"), candidate_globs=_budget_model_candidates(entry)) + for route in routes + ) + yield from _bedrock_request_forms( + bedrock_models=tuple(model for model in _bedrock_base_models() if entry in _bedrock_names(model)), + routes=routes, + ) + + +def _pattern_request_forms(pattern: str, routed_models: RoutedModels) -> Iterator[_RequestForm]: + yield _RequestForm(name_globs=(pattern,)) + yield from ( + _served_form(name) + for name in routed_models.served_names + if _GLOB_WILDCARD not in name and _glob_matches(pattern, _budget_model_candidates(name)) + ) + routes: Final = _wildcard_routes(routed_models) + yield from (_RequestForm(name_globs=(route, pattern), last_segment_open=True) for route in routes) + if "/" not in pattern: + yield from (_RequestForm(name_globs=(route, f"*/{pattern}"), candidate_globs=(pattern,)) for route in routes) + yield from ( + _narrowed_to(pattern=pattern, form=form) + for form in _bedrock_request_forms(bedrock_models=_bedrock_base_models(), routes=routes) + if _resolves(pattern=pattern, form=form) + ) + + +def _narrowed_to(pattern: str, form: _RequestForm) -> _RequestForm: + if any(_globs_share_a_name(pattern, candidate) for candidate in form.candidate_globs): + return form + return _RequestForm(name_globs=(*form.name_globs, pattern), candidate_globs=form.candidate_globs) + + +def _bedrock_request_forms(bedrock_models: Iterable[str], routes: Sequence[str]) -> Iterator[_RequestForm]: + """Wildcard-routed request names that resolve to a Bedrock model: its id, under a cross-region prefix or with a + context or throughput suffix, as the whole name or after any routing prefix, region path or ARN.""" + forms: Final = itertools.chain.from_iterable( + _bedrock_model_request_forms(model=model, routes=routes) for model in bedrock_models + ) + yield from (form for form in forms if _globs_share_a_name(*form.name_globs)) + + +def _bedrock_model_request_forms(model: str, routes: Sequence[str]) -> Iterator[_RequestForm]: + names: Final = _bedrock_names(model) + return ( + _RequestForm(name_globs=(route, f"{prefix}{last_segment}"), candidate_globs=(last_segment, *names)) + for route, (last_segment, prefix) in itertools.product( + routes, itertools.product(_bedrock_last_segments(model), ("", "*/")) + ) + ) + + +def _bedrock_last_segments(model: str) -> tuple[str, ...]: + regioned: Final = (model, *(f"{region}.{model}" for region in get_bedrock_cross_region_inference_regions())) + _, version_separator, version = model.rpartition(":") + suffixes: Final = ("", "[*]", *((":*k",) if version_separator and version.isdigit() else ())) + return tuple(f"{name}{suffix}" for name, suffix in itertools.product(regioned, suffixes)) + + +def _model_cost() -> Mapping[str, object]: + return litellm.model_cost + + +def _bedrock_base_models() -> tuple[str, ...]: + return tuple( + model + for model, cost_entry in _model_cost().items() + if _priced_as_bedrock(cost_entry) and get_bedrock_base_model(model) == model + ) + + +def _names_governed_by(entry: str, entry_members: frozenset[str], routed_models: RoutedModels) -> Iterator[str]: + yield entry + yield from (f"{prefix}{entry}" for prefix in _PROVIDER_PREFIXES) + yield from ( + name + for name in routed_models.served_names + if _GLOB_WILDCARD not in name and entry in _budget_model_candidates(name) + ) + yield from (member for member in entry_members if _GLOB_WILDCARD not in member) + + +def _ranks_ahead_of_group(entry: str, entry_members: frozenset[str], group: str, group_members: Iterable[str]) -> bool: + return any( + _GLOB_WILDCARD in member + or entry in _budget_model_candidates(member) + or (member in entry_members and entry < group) + for member in group_members + ) + + +def _usable_entry(entry: str, raw_budget_config: object) -> ResolvedModelBudget | None: + if (budget_config := _usable_budget_config(raw_budget_config)) is None: + # An entry that will not validate cannot be keyed, so it cannot be + # enforced or incremented. Skip to the next candidate rather than + # raising: raising would abort every other scope's increment and turn + # a config typo into a 500, and stopping here would let one malformed + # specific entry disable a perfectly good bare-family budget beside + # it. The candidate chain already falls through an ABSENT entry, and + # an unparseable one is indistinguishable from absent to enforcement. + # `validate_model_max_budget` rejects these on the write path, so + # reaching here means config.yaml or a direct DB edit. + verbose_proxy_logger.warning( + "Ignoring unusable model_max_budget entry for %s; it cannot be enforced or tracked", + entry, + ) + return None + return ResolvedModelBudget(budget_model=entry, budget_config=budget_config) + + +def team_model_budget_entries( + model: str, team_model_max_budget: Mapping[str, object], key_model_max_budget: Mapping[str, object] | None +) -> Mapping[str, object]: + """The team entries that budget `model` for this key. + + A key entry named after `model` that spend-gates it replaces the team entry named after it, so that model is + gated on and billed to the key's entry. The team's pattern and group entries still apply beside it. + """ + exact_candidates: Final = _budget_model_candidates(model) + if not _key_entry_replaces_team_entry( + model=model, exact_candidates=exact_candidates, key_model_max_budget=key_model_max_budget + ): + return team_model_max_budget + return MappingProxyType( + {entry: config for entry, config in team_model_max_budget.items() if entry not in exact_candidates} + ) + + +def _key_entry_replaces_team_entry( + model: str, exact_candidates: tuple[str, ...], key_model_max_budget: Mapping[str, object] | None +) -> bool: if not key_model_max_budget: - return True - resolved: Final = resolve_model_budget(model=model, model_max_budget=key_model_max_budget) - return resolved is None or not _spend_gated(resolved.budget_config) + return False + resolved: Final = resolve_model_budget( + model=model, + model_max_budget=MappingProxyType( + { + candidate: key_model_max_budget[candidate] + for candidate in exact_candidates + if candidate in key_model_max_budget + } + ), + model_access_groups=lambda _model, _team_id: frozenset(), + ) + return resolved is not None and _spend_gated(resolved.budget_config) def _spend_gated(budget_config: BudgetConfig) -> bool: return budget_config.max_budget is not None and budget_config.max_budget >= 0 +def entry_names_model(entry: str, model: str) -> bool: + """Whether `entry` budgets `model` under the model's own name, rather than as a pattern or an access group.""" + return entry in _budget_model_candidates(model) + + def _budget_model_candidates(model: str) -> tuple[str, ...]: """Names a budget may be configured under for a request on `model`, most specific first. @@ -189,13 +582,18 @@ def _bedrock_candidates(model: str) -> tuple[str, ...]: of those would produce a garbage candidate. """ base_model: Final = get_bedrock_base_model(model) - cost_entry: Final = litellm.model_cost.get(base_model) - if not isinstance(cost_entry, dict) or not str(cost_entry.get("litellm_provider", "")).startswith("bedrock"): - return () + return _bedrock_names(base_model) if _priced_as_bedrock(_model_cost().get(base_model)) else () + + +def _bedrock_names(base_model: str) -> tuple[str, ...]: _, _, without_vendor = base_model.partition(".") return (base_model, without_vendor) if without_vendor else (base_model,) +def _priced_as_bedrock(cost_entry: object) -> bool: + return str(_metadata_value(cost_entry, "litellm_provider") or "").startswith("bedrock") + + @dataclass(frozen=True, slots=True) class _WindowReading: current_spend: float @@ -208,6 +606,7 @@ async def build_model_max_budget_usage( model_max_budget: Mapping[str, object] | None, cache: DualCache | None, now: datetime | None = None, + shared_windows: SharedWindows = PROXY_SHARED_WINDOWS, ) -> dict[str, dict[str, object]]: """Current-window spend per configured budget model, as `/key/info` and `/user/info` report it. @@ -231,17 +630,18 @@ async def build_model_max_budget_usage( ) if not budgets: return {} + spend_keys: Final = tuple( + model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=budget_model, + budget_duration=budget_config.budget_duration, + ) + for budget_model, budget_config in budgets + ) windows: Final = await _read_windows( cache=cache, - spend_keys=tuple( - model_budget_spend_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=budget_model, - budget_duration=budget_config.budget_duration, - ) - for budget_model, budget_config in budgets - ), + spend_keys=spend_keys, start_time_keys=tuple( model_budget_start_time_cache_key( entity_type=entity_type, @@ -253,21 +653,57 @@ async def build_model_max_budget_usage( ), ) current_time: Final = (now or datetime.now(timezone.utc)).timestamp() + shared: Final = tuple( + [ + ( + await _open_shared_window( + shared_windows=shared_windows, + spend_key=spend_key, + budget_duration=budget_config.budget_duration, + now=current_time, + ) + if cache.redis_cache is None + else None + ) + for spend_key, (_, budget_config) in zip(spend_keys, budgets, strict=True) + ] + ) + readings: Final = tuple( + _fuller_reading(local=window, shared=shared_window) + for window, shared_window in zip(windows, shared, strict=True) + ) return { budget_model: { - "current_spend": round(window.current_spend, 4), + "current_spend": round(reading.current_spend, 4), "budget_limit": budget_config.max_budget, "time_period": budget_config.budget_duration, "budget_reset_at": _window_reset_at( - started_at=window.started_at, + started_at=reading.started_at, budget_duration=budget_config.budget_duration, now=current_time, ), } - for (budget_model, budget_config), window in zip(budgets, windows, strict=True) + for (budget_model, budget_config), reading in zip(budgets, readings, strict=True) } +async def _open_shared_window( + shared_windows: SharedWindows, spend_key: str, budget_duration: str | None, now: float +) -> SharedWindow | None: + if budget_duration is None: + return None + window: Final = await shared_windows.read(spend_key=spend_key, budget_duration=budget_duration) + if window is None or not window.is_open(budget_duration=budget_duration, now=now): + return None + return window + + +def _fuller_reading(local: _WindowReading, shared: SharedWindow | None) -> _WindowReading: + if shared is None or shared.spend <= local.current_spend: + return local + return _WindowReading(current_spend=shared.spend, started_at=shared.started_at) + + @with_service_target("model_budgets") async def _read_windows( cache: DualCache, spend_keys: Sequence[str], start_time_keys: Sequence[str] @@ -307,6 +743,10 @@ def _window_reset_at(started_at: float | None, budget_duration: str | None, now: return datetime.fromtimestamp(ends_at, tz=timezone.utc).isoformat() +def is_enforced_budget_entry(raw_budget_config: object) -> bool: + return _usable_budget_config(raw_budget_config) is not None + + def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None: try: budget_config: Final = BudgetConfig.model_validate(raw_budget_config) @@ -337,18 +777,44 @@ def _as_float(value: object) -> float | None: def _resolve_entity_model_budgets( model: str, entity_budgets: Iterable[tuple[Litellm_EntityType, str | None, object]], + model_access_groups: ModelAccessGroups, + team_id: str | None, ) -> tuple[tuple[Litellm_EntityType, str, ResolvedModelBudget], ...]: """Drop the scopes that do not budget `model`, keeping only what can be incremented.""" return tuple( (entity_type, entity_id, resolved) for entity_type, entity_id, model_max_budget in entity_budgets if entity_id is not None and isinstance(model_max_budget, Mapping) and model_max_budget - for resolved in (resolve_model_budget(model=model, model_max_budget=model_max_budget),) + for resolved in ( + resolve_model_budget( + model=model, + model_max_budget=model_max_budget, + model_access_groups=model_access_groups, + team_id=team_id, + ), + ) if resolved is not None and resolved.budget_config.budget_duration is not None ) +def _requested_model_name( + metadata: Mapping[str, object], standard_logging_payload: StandardLoggingPayload +) -> str | None: + """The model name the request was gated on: a team-scoped deployment's public name, since the router renames + its model group to a unique internal one, else the model group, else the deployment's model.""" + team_public_model_name: Final = _metadata_value(metadata.get("model_info"), "team_public_model_name") + if ( + isinstance(team_public_model_name, str) + and team_public_model_name + and _GLOB_WILDCARD not in team_public_model_name + ): + return team_public_model_name + return standard_logging_payload.get("model_group") or standard_logging_payload.get("model") + + def _metadata_value(metadata: object, key: str) -> object | None: + if not isinstance(metadata, Mapping): + return None metadata_mapping: Final[Mapping[str, object]] = cast( # cast-ok: callback metadata has string keys Mapping[str, object], metadata ) @@ -367,9 +833,17 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): Example: key=$LITELLM_API_KEY, model=gpt-4o, max_budget=100, time_period=1d """ - def __init__(self, dual_cache: DualCache, clock: Callable[[], float] = time.time): + def __init__( + self, + dual_cache: DualCache, + clock: Callable[[], float] = time.time, + model_access_groups: ModelAccessGroups = proxy_model_access_groups, + shared_windows: SharedWindows = PROXY_SHARED_WINDOWS, + ): self.dual_cache = dual_cache self._clock = clock + self._model_access_groups = model_access_groups + self._shared_windows = shared_windows self.redis_increment_operation_queue = [] self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() @@ -393,6 +867,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return await self._is_entity_within_model_budget( entity_type=Litellm_EntityType.KEY, entity_id=user_api_key_dict.token, + team_id=user_api_key_dict.team_id, model_max_budget=user_api_key_dict.model_max_budget, model=model, exceeded_message=( @@ -422,6 +897,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): user_id: str, user_model_max_budget: Mapping[str, object], model: str, + team_id: str | None = None, ) -> bool: """ Check if the internal user is within the model budget @@ -432,6 +908,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return await self._is_entity_within_model_budget( entity_type=Litellm_EntityType.USER, entity_id=user_id, + team_id=team_id, model_max_budget=user_model_max_budget, model=model, exceeded_message=f"LiteLLM User: {user_id}, exceeded budget for model={model}", @@ -443,6 +920,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): end_user_id: str, end_user_model_max_budget: Mapping[str, object], model: str, + team_id: str | None = None, ) -> bool: """ Check if the end_user is within the model budget @@ -453,6 +931,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return await self._is_entity_within_model_budget( entity_type=Litellm_EntityType.END_USER, entity_id=end_user_id, + team_id=team_id, model_max_budget=end_user_model_max_budget, model=model, exceeded_message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", @@ -467,18 +946,19 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): model: str, ) -> bool: """ - Check if the team is within the model budget, unless the key's own - `model_max_budget` overrides it for `model` + Check if the team is within the model budget, leaving out a team entry named after `model` that the key's + own `model_max_budget` replaces Raises: BudgetExceededError: If the team has exceeded the model budget """ - if not team_model_budget_applies(model=model, key_model_max_budget=key_model_max_budget): - return True return await self._is_entity_within_model_budget( entity_type=Litellm_EntityType.TEAM, entity_id=team_id, - model_max_budget=team_model_max_budget, + team_id=team_id, + model_max_budget=team_model_budget_entries( + model=model, team_model_max_budget=team_model_max_budget, key_model_max_budget=key_model_max_budget + ), model=model, exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}", ) @@ -494,6 +974,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return await self._is_entity_within_model_budget( entity_type=Litellm_EntityType.TEAM_MEMBER, entity_id=team_member_budget_entity_id(user_id=user_id, team_id=team_id), + team_id=team_id, model_max_budget=team_member_model_max_budget, model=model, exceeded_message=( @@ -505,13 +986,19 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): self, entity_type: Litellm_EntityType, entity_id: str | None, + team_id: str | None, model_max_budget: Mapping[str, object] | None, model: str, exceeded_message: str, ) -> bool: if not model_max_budget: return True - resolved: Final = resolve_model_budget(model=model, model_max_budget=model_max_budget) + resolved: Final = resolve_model_budget( + model=model, + model_max_budget=model_max_budget, + model_access_groups=self._model_access_groups, + team_id=team_id, + ) if resolved is None: verbose_proxy_logger.debug("Model %s not found in %s model_max_budget", model, entity_type.value) return True @@ -520,15 +1007,26 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): if max_budget is None or max_budget < 0: return True - current_spend: Final = await self._get_spend_for_model_budget( + local_spend: Final = await self._get_spend_for_model_budget( entity_type=entity_type, entity_id=entity_id, model=model, resolved=resolved, ) + current_spend: Final = ( + local_spend + if local_spend >= max_budget + else await self._with_shared_spend( + entity_type=entity_type, entity_id=entity_id, resolved=resolved, local_spend=local_spend + ) + ) if current_spend >= max_budget: raise litellm.BudgetExceededError( - message=exceeded_message, + message=( + exceeded_message + if resolved.budget_model in _budget_model_candidates(model) + else f"{exceeded_message}, budget_model={resolved.budget_model}" + ), current_cost=current_spend, max_budget=max_budget, entity_type=entity_type.value, @@ -566,6 +1064,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return current_spend return current_spend + _as_spend(await self._cached_spend(legacy_spend_key)) + async def _with_shared_spend( + self, + entity_type: Litellm_EntityType, + entity_id: str | None, + resolved: ResolvedModelBudget, + local_spend: float, + ) -> float: + """Without Redis each replica's counter holds only its own share, so the window every replica adds to + in the database is the floor.""" + if not self._shares_windows_in_database(entity_type): + return local_spend + shared: Final = await _open_shared_window( + shared_windows=self._shared_windows, + spend_key=model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=resolved.budget_config.budget_duration, + ), + budget_duration=resolved.budget_config.budget_duration, + now=self._clock(), + ) + return local_spend if shared is None else max(local_spend, shared.spend) + async def _cached_spend(self, spend_key: str) -> float | None: redis_cache: Final = self.dual_cache.redis_cache if redis_cache is None: @@ -604,13 +1126,10 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) payload_metadata: Final = standard_logging_payload.get("metadata") or {} - # Use model_group (the user-facing model alias, e.g. "gpt-4o") when - # available. The enforcement path receives the model name from - # request_data["model"] which is the model group alias, so the spend - # tracking cache key must resolve from the same name. Falling back to - # the deployment-level "model" field preserves behaviour for non-proxy - # or non-router deployments where model_group is None. - model: Final = standard_logging_payload.get("model_group") or standard_logging_payload.get("model") + model: Final = _requested_model_name( + metadata=_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped + standard_logging_payload=standard_logging_payload, + ) if model is None: return @@ -624,6 +1143,10 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): payload_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped "user_api_key_team_id", ) + team_model_max_budget: Final = _metadata_value( + _metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped + "user_api_key_team_model_max_budget", + ) team_member_entity_id: Final = ( team_member_budget_entity_id(user_id=team_member_user_id, team_id=team_member_team_id) if team_member_user_id is not None and team_member_team_id is not None @@ -637,16 +1160,17 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ), ( Litellm_EntityType.TEAM, - payload_metadata.get("user_api_key_team_id"), + team_member_team_id, ( - _metadata.get("user_api_key_team_model_max_budget") - if team_model_budget_applies( + team_model_budget_entries( model=model, + team_model_max_budget=team_model_max_budget, key_model_max_budget=( key_model_max_budget if isinstance(key_model_max_budget, Mapping) else None ), ) - else None + if isinstance(team_model_max_budget, Mapping) + else team_model_max_budget ), ), ( @@ -669,7 +1193,12 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ), ) - resolved_budgets: Final = _resolve_entity_model_budgets(model=model, entity_budgets=entity_budgets) + resolved_budgets: Final = _resolve_entity_model_budgets( + model=model, + entity_budgets=entity_budgets, + model_access_groups=self._model_access_groups, + team_id=team_member_team_id, + ) if not resolved_budgets: verbose_proxy_logger.debug( "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: " @@ -723,17 +1252,114 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ttl_seconds=duration_in_seconds(budget_duration), ): return + start_time_key: Final = model_budget_start_time_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=budget_duration, + ) + shares_window: Final = self._shares_windows_in_database(entity_type) + if shares_window: + await self._join_shared_window( + spend_key=spend_key, start_time_key=start_time_key, budget_duration=budget_duration + ) await self._increment_spend_for_key( budget_config=resolved.budget_config, spend_key=spend_key, - start_time_key=model_budget_start_time_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=resolved.budget_model, - budget_duration=budget_duration, - ), + start_time_key=start_time_key, response_cost=response_cost, ) + if shares_window: + await self._share_spend( + spend_key=spend_key, + start_time_key=start_time_key, + budget_duration=budget_duration, + response_cost=response_cost, + ) + + def _shares_windows_in_database(self, entity_type: Litellm_EntityType) -> bool: + """End-user ids are whatever the client sends, so their windows stay on each replica rather than adding a + database row per id.""" + return self.dual_cache.redis_cache is None and entity_type not in _PER_REPLICA_SCOPES + + async def _join_shared_window(self, spend_key: str, start_time_key: str, budget_duration: str) -> None: + """Without Redis, start this replica's window where the database says every replica's window starts. + + Every replica then sends the same window start with its spend, so no flush reads as a newer window that + replaces the spend the others already added. While the database cannot say, this replica keeps its spend + unshared and sends it once a later charge joins the shared window. + """ + now: Final = self._clock() + local: Final = _local_windows(self.dual_cache) + unshared_key: Final = _unshared_spend_key(start_time_key) + if local.get_cache(key=unshared_key) is None and _open_local_window_start( + local=local, start_time_key=start_time_key, budget_duration=budget_duration, now=now + ): + return + started_at: Final = await self._shared_windows.open( + spend_key=spend_key, budget_duration=budget_duration, now=now + ) + unshared_spend: Final = _as_float(local.get_cache(key=unshared_key)) + local_window_open: Final = _open_local_window_start( + local=local, start_time_key=start_time_key, budget_duration=budget_duration, now=now + ) + if local_window_open and unshared_spend is None: + return + if not local_window_open: + _start_local_window( + local=local, + spend_key=spend_key, + start_time_key=start_time_key, + started_at=now if started_at is None else started_at, + budget_duration=budget_duration, + now=now, + spend=0.0, + ) + if started_at is None: + local.set_cache(key=unshared_key, value=0.0, ttl=budget_window_seconds(now, budget_duration)) + else: + local.delete_cache(key=unshared_key) + return + if started_at is None or unshared_spend is None: + return + _start_local_window( + local=local, + spend_key=spend_key, + start_time_key=start_time_key, + started_at=started_at, + budget_duration=budget_duration, + now=now, + spend=_as_float(local.get_cache(key=spend_key)) or 0.0, + ) + local.delete_cache(key=unshared_key) + await self._shared_windows.add( + spend_key=spend_key, budget_duration=budget_duration, started_at=started_at, spend=unshared_spend + ) + + async def _share_spend( + self, spend_key: str, start_time_key: str, budget_duration: str, response_cost: float + ) -> None: + now: Final = self._clock() + local: Final = _local_windows(self.dual_cache) + unshared_key: Final = _unshared_spend_key(start_time_key) + unshared_spend: Final = _as_float(local.get_cache(key=unshared_key)) + if unshared_spend is not None: + local.set_cache( + key=unshared_key, + value=unshared_spend + response_cost, + ttl=budget_window_seconds(now, budget_duration), + ) + return + await self._shared_windows.add( + spend_key=spend_key, + budget_duration=budget_duration, + started_at=await self._get_or_set_budget_start_time( + start_time_key=start_time_key, + current_time=now, + ttl_seconds=budget_window_seconds(now, budget_duration), + ), + spend=response_cost, + ) async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) @@ -743,4 +1369,33 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return polls == 1 +def _local_windows(dual_cache: DualCache) -> LocalTTLCache: + return dual_cache.in_memory_cache + + +def _open_local_window_start(local: LocalTTLCache, start_time_key: str, budget_duration: str, now: float) -> bool: + started_at: Final = _as_float(local.get_cache(key=start_time_key)) + return started_at is not None and now - started_at <= budget_window_seconds(started_at, budget_duration) + + +def _start_local_window( + local: LocalTTLCache, + spend_key: str, + start_time_key: str, + started_at: float, + budget_duration: str, + now: float, + spend: float, +) -> None: + ttl: Final = started_at + budget_window_seconds(started_at, budget_duration) - now + local.delete_cache(key=spend_key) + local.delete_cache(key=start_time_key) + local.set_cache(key=spend_key, value=spend, ttl=ttl) + local.set_cache(key=start_time_key, value=started_at, ttl=ttl) + + +def _unshared_spend_key(start_time_key: str) -> str: + return f"{start_time_key}:unshared" + + PROXY_VirtualKeyModelMaxBudgetLimiter = _PROXY_VirtualKeyModelMaxBudgetLimiter diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 567b8452ae2..e44d9304ac1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -45,11 +45,14 @@ from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.cost_calculator import get_usage_object from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, + split_media_for_tpm, ) from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import estimate_image_tokens, offload_token_count from litellm.llms.azure.ptu_capacity import normalized_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( @@ -98,6 +101,7 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + from litellm.router import Router from litellm.types.agents import AgentResponse from litellm.types.caching import RedisPipelineIncrementOperation @@ -135,6 +139,18 @@ class RateLimitedModel: return requested_limit if requested_limit is not None else limits.get(self.group) +def _resolve_media_models(router: "Router | None", model: str, team_id: str | None) -> tuple[str, ...]: + if router is None: + return (model,) + return tuple(dict.fromkeys(router.resolved_litellm_models(model, team_id, use_base_model=True))) or (model,) + + +def _resolve_media_models_via_proxy_router(model: str, team_id: str | None) -> tuple[str, ...]: + from litellm.proxy.proxy_server import llm_router + + return _resolve_media_models(llm_router, model, team_id) + + def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None: from litellm.proxy.proxy_server import llm_router @@ -862,6 +878,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): time_provider: Callable[[], datetime] | None = None, tag_rate_limit_resolver: TagRateLimitResolver = resolve_tag_rate_limits_from_db, model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router, + media_model_resolver: Callable[[str, str | None], tuple[str, ...]] = _resolve_media_models_via_proxy_router, ptu_team_ceiling_resolver: Callable[ [str, str], PTUTeamCeiling | None ] = _resolve_ptu_team_ceiling_via_proxy_router, @@ -872,6 +889,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._time_provider = time_provider or datetime.now self._tag_rate_limit_resolver = tag_rate_limit_resolver self._model_group_resolver = model_group_resolver + self._media_model_resolver = media_model_resolver self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver self._fail_closed_resolver = fail_closed_resolver self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver @@ -1133,11 +1151,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _estimate_tokens_for_request( self, - data: dict, + data: object, model: str | None = None, min_configured_tpm_limit: int | None = None, call_type: str | None = None, configured_output_tokens: int | None = None, + team_id: str | None = None, ) -> int: """ Estimate total tokens this request will consume so we can reserve them @@ -1164,6 +1183,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=team_id, + media_model_resolver=self._media_model_resolver, ) total_estimated: Final = estimated_input_tokens + max_tokens_estimate @@ -1184,6 +1205,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): call_type: str | None, configured_output_tokens: int | None, raw_estimate: int, + team_id: str | None = None, ) -> int: """The team PTU ceiling counts Azure normalized tokens, so its reservation weighs the output budget the way the ceiling does instead of the raw sum the other scopes reserve.""" @@ -1194,6 +1216,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=team_id, + media_model_resolver=self._media_model_resolver, ) normalized: Final = normalized_tokens( ceiling, prompt_tokens=estimated_input_tokens, completion_tokens=max_tokens_estimate @@ -1206,6 +1230,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit: int | None = None, call_type: str | None = None, configured_output_tokens: int | None = None, + include_media: bool = True, + team_id: str | None = None, + media_model_resolver: Callable[[str, str | None], tuple[str, ...]] | None = None, ) -> tuple[int, int]: """ Estimate input tokens and output (max_tokens) budget separately, so @@ -1229,10 +1256,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): floor entirely, so the reservation reflects what this tenant's model actually emits rather than one constant shared by every tenant. """ - if not isinstance(data, dict): + request_data: Final = as_str_mapping(data) + if not isinstance(data, dict) or request_data is None: return 0, 0 + request_model: Final = request_data.get("model") translated_data: Final = self._translate_google_genai_native_request(data, call_type) - estimable_data: Final = translated_data if translated_data is not None else data + estimable_data: Final = translated_data if translated_data is not None else request_data selected_fields: Final[tuple[object | None, object | None, object | None]] = ( (None, None, estimable_data.get("input")) if call_type in RESPONSES_API_CALL_TYPES or call_type in EMBEDDING_API_CALL_TYPES @@ -1249,10 +1278,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) messages, prompt, raw_input = selected_fields - input_text: Final = ( + original_input: Final = ( billable_input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input ) + media_source: Final = messages if messages is not None else original_input + projected, media = ( + split_media_for_tpm(media_source) + if include_media and (isinstance(messages, list) or call_type in RESPONSES_API_CALL_TYPES) + else (original_input, ()) + ) + input_text: Final = projected if call_type in RESPONSES_API_CALL_TYPES else original_input + fallback_model: Final = request_model if isinstance(request_model, str) else None + resolved_models: Final = ( + media_model_resolver(fallback_model or "", team_id) + if media and media_model_resolver is not None + else (fallback_model or "",) + ) + media_tokens: Final = sum( + self._estimate_media_block_tokens(block, resolved_models, fallback_model) for block in media + ) + total_chars: Final = ( len(get_str_from_messages(messages)) if isinstance(messages, list) and messages @@ -1267,7 +1313,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else 0 ) - estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 + estimated_input_tokens: Final = ( + max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0 + ) + media_tokens explicit_max_tokens: Final = self._get_explicit_output_cap(data, call_type) generates_no_output: Final = ( @@ -1283,7 +1331,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): max_tokens_estimate: Final = ( 0 if generates_no_output - or (explicit_max_tokens is None and total_chars == 0 and configured_output_tokens is None) + or (explicit_max_tokens is None and estimated_input_tokens == 0 and configured_output_tokens is None) else explicit_max_tokens if explicit_max_tokens is not None else configured_output_tokens @@ -1293,6 +1341,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return estimated_input_tokens, max_tokens_estimate * self.get_output_candidate_count(data, call_type) + @staticmethod + def _estimate_media_block_tokens( + block: Mapping[str, object], models: tuple[str, ...], fallback_model: str | None + ) -> int: + if block.get("type") == "input_audio": + audio: Final = _callback_kwargs_mapping(block.get("input_audio")) + data: Final = audio.get("data") + return estimate_audio_tokens(data if isinstance(data, str) else None, models, fallback_model=fallback_model) + image: Final = _callback_kwargs_mapping(block.get("image_url")) + url: Final = image.get("url") + detail: Final = image.get("detail") + return estimate_image_tokens( + models, + url if isinstance(url, str) else "", + detail if isinstance(detail, str) else "auto", + fallback_model=fallback_model, + ) + def _is_redis_cluster(self) -> bool: """ Check if the dual cache is using Redis cluster. @@ -3958,7 +4024,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): except Exception: # noqa: BLE001 # tokenizer failures degrade to the cheap estimate if call_type in RERANK_API_CALL_TYPES and isinstance(selected_text, str): return max(0, len(selected_text) // DEFAULT_CHARS_PER_TOKEN) - estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data, call_type=call_type) + estimated_input_tokens, _ = self._estimate_input_and_output_tokens( + data=data, call_type=call_type, include_media=False + ) return estimated_input_tokens + audio_token_estimate async def _reserve_project_io_tokens_or_raise( @@ -4003,6 +4071,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data=data, min_configured_tpm_limit=min_configured_otpm_limit, call_type=call_type, + include_media=False, ) raw_estimated_input_tokens: Final = await offload_token_count(self._estimate_precise_input_tokens)( data=data, model=requested_model, call_type=call_type @@ -4360,12 +4429,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Post-call reconciliation refunds the over-reservation # delta when actual usage comes in below the floor. estimated_tokens: Final = max( - self._estimate_tokens_for_request( + await offload_token_count(self._estimate_tokens_for_request)( data=data, model=requested_model, min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, + team_id=user_api_key_dict.team_id, ), 1, ) @@ -4383,13 +4453,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) capped_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data) - ptu_estimated_tokens: Final = self._estimate_ptu_tokens_for_request( + ptu_estimated_tokens: Final = await offload_token_count(self._estimate_ptu_tokens_for_request)( ceiling=stash.ptu_ceiling, data=capped_request, min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, raw_estimate=estimated_tokens, + team_id=user_api_key_dict.team_id, ) if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit: if not reuses_request_slot: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b22049690a1..a6a60fec2ac 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -119,6 +119,9 @@ _SESSION_ID_VALUE_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") _SHA256_HEX_RE: Final = re.compile(r"^[0-9a-f]{64}$") +SCHEDULER_PRIORITY_HEADER: Final = "x-litellm-priority" +ROUTER_SCHEDULED_ROUTE_TYPES: Final = frozenset({"acompletion", "atext_completion"}) + # W3C Trace Context traceparent header: https://www.w3.org/TR/trace-context/ # e.g. "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" _TRACEPARENT_RE: Final = re.compile(r"^[0-9a-f]{2}-([0-9a-f]{32})-[0-9a-f]{16}-[0-9a-f]{2}$", re.IGNORECASE) @@ -689,7 +692,7 @@ def _metadata_string_tags(metadata_objects: Iterable[Mapping[str, object]]) -> t def _strip_untrusted_request_header_controls( - headers: Any, + headers: object, *, allow_client_message_redaction_opt_out: bool = False, ) -> None: @@ -1693,7 +1696,7 @@ class LiteLLMProxyRequestSetup: @staticmethod def add_internal_user_from_user_mapping( - general_settings: dict | None, + general_settings: Mapping[str, object] | None, user_api_key_dict: UserAPIKeyAuth, headers: dict, ) -> UserAPIKeyAuth: @@ -1712,7 +1715,7 @@ class LiteLLMProxyRequestSetup: return user_api_key_dict @staticmethod - def get_user_from_headers(headers: dict, general_settings: dict | None = None) -> str | None: + def get_user_from_headers(headers: dict, general_settings: Mapping[str, object] | None = None) -> str | None: """ Get the user from the specified header if `general_settings.user_header_name` is set. """ @@ -1733,7 +1736,9 @@ class LiteLLMProxyRequestSetup: return user @staticmethod - def get_openai_org_id_from_headers(headers: dict, general_settings: dict | None = None) -> str | None: + def get_openai_org_id_from_headers( + headers: dict, general_settings: Mapping[str, object] | None = None + ) -> str | None: """ Get the OpenAI Org ID from the headers. """ @@ -1824,7 +1829,7 @@ class LiteLLMProxyRequestSetup: headers: dict, request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - general_settings: dict[str, object] | None = None, + general_settings: Mapping[str, object] | None = None, ) -> LitellmDataForBackendLLMCall: """ - Adds user from headers @@ -3285,7 +3290,7 @@ def _per_user_credential_names_for_groups( GITHUB_COPILOT_AUTH_TYPE_KEY, GITHUB_COPILOT_PER_USER_AUTH_TYPE, ) - from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model + from litellm.router_strategy.auto_router.auto_router_model_naming import classify_strategy_router_model names: Final[list[str]] = [] # mutable-ok: accumulates one name per per-user deployment for group in model_groups: diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 7a843aecc0c..3dceab6b7f7 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -60,16 +60,16 @@ from litellm.repositories.table_repositories import PrismaTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + StrategyRouterDependencyRole, + classify_strategy_router_model, + strategy_router_dependencies, +) from litellm.router_strategy.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.request_models import ( AutoRouterRoutingTestRequest, RequestComplexityRouterConfig, ) -from litellm.router_utils.auto_router_model_naming import ( - StrategyRouterDependencyRole, - classify_strategy_router_model, - strategy_router_dependencies, -) from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.management_endpoints.auto_router_endpoints import ( SHADOW_EVAL_TURN_VALVE, @@ -231,9 +231,12 @@ def _shadow_eval_attempts(prisma_client: "PrismaClient") -> _ShadowEvalAttemptTa return _ShadowEvalAttemptRepository(prisma_client).table +_RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(strict=True)) + + async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]: async with db_span("auto_router_report_query", sql_relation(query)): - return await prisma_client.db.query_raw(query, *args) + return _RAW_ROWS.validate_python(await prisma_client.db.query_raw(query, *args)) async def _execute_raw(prisma_client: "PrismaClient", query: str, *args: object) -> int: @@ -408,7 +411,7 @@ async def validate_complexity_router_config( auto_router_environment_reference_error(data.complexity_router_config, user_api_key_dict) ) - from litellm.router_utils.auto_router_model_naming import ( + from litellm.router_strategy.auto_router.auto_router_model_naming import ( validate_complexity_router_config_write, ) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index f948a49441a..66f386b792d 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -3,6 +3,7 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass, replace from datetime import date, datetime, timedelta +from itertools import groupby from types import MappingProxyType from typing import Final, Literal, NoReturn, Protocol @@ -28,6 +29,10 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendData, DailySpendMetadata, GroupedData, + GuardrailKeyMetric, + GuardrailMetadata, + GuardrailMetric, + GuardrailSpendMetrics, KeyMetadata, KeyMetricWithMetadata, MetricWithMetadata, @@ -39,6 +44,7 @@ from litellm.types.repositories.daily_activity import ( DailyActivityTable, EntityRollupRow, GroupingSetsRow, + GuardrailSpendRow, KeyMetadataRow, RollupMetricsRow, SpendLogsWindow, @@ -1013,17 +1019,68 @@ def _fold_entity_rollups_sync( ) +def _guardrail_metric(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> GuardrailMetric: + return GuardrailMetric( + metrics=GuardrailSpendMetrics( + spend=sum(row.spend for row in rows), + untracked_units=sum(row.untracked_units for row in rows), + ), + metadata=GuardrailMetadata(guardrail_name=rows[0].guardrail_name), + api_key_breakdown={ + row.api_key: GuardrailKeyMetric( + metrics=GuardrailSpendMetrics(spend=row.spend, untracked_units=row.untracked_units), + metadata=day.breakdown.api_keys[row.api_key].metadata, + ) + for row in rows + if row.api_key in day.breakdown.api_keys + }, + ) + + +def _day_with_guardrails(day: DailySpendData, rows: Sequence[GuardrailSpendRow]) -> DailySpendData: + by_guardrail: Final = groupby(sorted(rows, key=lambda row: row.guardrail_id), key=lambda row: row.guardrail_id) + guardrails: Final = {guardrail_id: _guardrail_metric(day, tuple(group)) for guardrail_id, group in by_guardrail} + return day.model_copy(update={"breakdown": day.breakdown.model_copy(update={"guardrails": guardrails})}) + + +def _with_guardrail_breakdown( + results: Sequence[DailySpendData], rows: Sequence[GuardrailSpendRow] +) -> tuple[DailySpendData, ...]: + rows_by_date: Final = MappingProxyType( + {day: tuple(group) for day, group in groupby(sorted(rows, key=lambda row: row.date), key=lambda row: row.date)} + ) + return tuple(_day_with_guardrails(day, rows_by_date.get(day.date.strftime("%Y-%m-%d"), ())) for day in results) + + +async def _guardrail_spend_or_empty( + repository: DailyActivityRepository, scope: DailyActivityScope +) -> tuple[GuardrailSpendRow, ...]: + from prisma.errors import PrismaError + + try: + return await repository.guardrail_spend(scope) + except PrismaError as exc: + verbose_proxy_logger.warning("Returning daily activity without its guardrail breakdown: %s", exc) + return () + + async def get_daily_activity_aggregated( repository: DailyActivityRepository, scope: DailyActivityScope, *, entity_metadata_field: Mapping[str, dict[str, object]] | None = None, include_entity_breakdown: bool = False, + include_guardrail_breakdown: bool = False, api_key_limit: int = constants.USAGE_TOP_API_KEYS_DEFAULT, ) -> SpendAnalyticsPaginatedResponse: try: - aggregated_rows: Final = await repository.aggregated( - scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit + aggregated_rows, guardrail_rows = await asyncio.gather( + repository.aggregated( + scope, include_entity_breakdown=include_entity_breakdown, api_key_limit=api_key_limit + ), + _guardrail_spend_or_empty(repository, scope) + if include_guardrail_breakdown and not scope.model + else asyncio.sleep(0, result=()), ) records: Final = aggregated_rows.grouping_rows aggregated: Final = await _aggregate_grouping_sets_records( @@ -1058,8 +1115,13 @@ async def get_daily_activity_aggregated( api_key_metadata=entity_key_metadata, entity_metadata_field=entity_metadata_field, ) + results: Final = ( + list(await asyncio.to_thread(_with_guardrail_breakdown, aggregated["results"], guardrail_rows)) + if guardrail_rows + else aggregated["results"] + ) return SpendAnalyticsPaginatedResponse( - results=aggregated["results"], + results=results, metadata=DailySpendMetadata( total_spend=aggregated["totals"].spend, total_flat_cost=aggregated["totals"].flat_cost, diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 842873129e4..79e515944ce 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -300,6 +300,7 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver, resolved.scope, entity_metadata_field=resolved.entity_metadata, include_entity_breakdown=resolver.include_entity_breakdown, + include_guardrail_breakdown=resolver.include_guardrail_breakdown, api_key_limit=api_key_limit, ) if resolver.entity != "team": diff --git a/litellm/proxy/management_endpoints/daily_activity_scopes.py b/litellm/proxy/management_endpoints/daily_activity_scopes.py index bc373388a4e..77246f5dda6 100644 --- a/litellm/proxy/management_endpoints/daily_activity_scopes.py +++ b/litellm/proxy/management_endpoints/daily_activity_scopes.py @@ -56,6 +56,7 @@ class EntityScopeResolver: resolve: EntityScopeResolve include_entity_breakdown: bool operation_names: OperationNames + include_guardrail_breakdown: bool = False def _query_ids(value: str | None) -> tuple[str, ...] | None: @@ -364,6 +365,7 @@ USER_RESOLVER = EntityScopeResolver( query=_user_query, resolve=_resolve_user, include_entity_breakdown=False, + include_guardrail_breakdown=True, operation_names=MappingProxyType( { "aggregated": "get_user_daily_activity_aggregated", @@ -385,6 +387,7 @@ TEAM_RESOLVER = EntityScopeResolver( query=_team_query, resolve=_resolve_team, include_entity_breakdown=True, + include_guardrail_breakdown=True, operation_names=MappingProxyType( { "aggregated": "get_team_daily_activity_aggregated", @@ -423,6 +426,7 @@ ORGANIZATION_RESOLVER = EntityScopeResolver( query=_organization_query, resolve=_resolve_organization, include_entity_breakdown=True, + include_guardrail_breakdown=True, operation_names=MappingProxyType( { "aggregated": "get_organization_daily_activity_aggregated", diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d6bf9220551..f2f0bc5be0b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -499,7 +499,7 @@ async def new_user( - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) + - model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user. @@ -1702,7 +1702,7 @@ async def user_update( - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) + - model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index f7b72b263d4..a19ecb8a9e9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2328,7 +2328,7 @@ async def generate_key_fn( - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only. - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. + - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. @@ -2538,7 +2538,7 @@ async def generate_service_account_key_fn( - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - guardrails: Optional[List[str]] - List of active guardrails for the key - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. + - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. @@ -3855,7 +3855,7 @@ async def update_key_fn( - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) - spend: Optional[float] - Amount spent by key - max_budget: Optional[float] - Max budget for key - - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget. @@ -6303,7 +6303,7 @@ async def regenerate_key_fn( - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) - spend: Optional[float] - Amount spent by key - max_budget: Optional[float] - Max budget for key - - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 45006b14df6..8cd3559ba86 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -131,6 +131,23 @@ from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ModelTableRepository from litellm.repositories.team_repository import TeamRepository from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + GATED_AUTO_ROUTER_CAPABILITIES, + STRATEGY_ROUTER_PARAM_FIELDS, + GatedAutoRouterCapability, + capability_limit_violation, + carries_complexity_router_settings, + count_capability_routers, + gated_capabilities_of, + is_complexity_router_model, + validate_complexity_router_config_placement, + validate_complexity_router_config_write, + validate_strategy_router_model_write, +) +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( + is_mutable_tuned_candidate, + tuning_quota_violation, +) from litellm.router_strategy.complexity_router import ( DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, ClassificationRubric, @@ -144,20 +161,6 @@ from litellm.router_strategy.complexity_router import ( normalize_classification_prompt, ) from litellm.router_strategy.complexity_router.config import resolve_complexity_router_config_write -from litellm.router_utils.auto_router_model_naming import ( - GATED_AUTO_ROUTER_CAPABILITIES, - STRATEGY_ROUTER_PARAM_FIELDS, - GatedAutoRouterCapability, - capability_limit_violation, - carries_complexity_router_settings, - count_capability_routers, - gated_capabilities_of, - is_complexity_router_model, - validate_complexity_router_config_placement, - validate_complexity_router_config_write, - validate_strategy_router_model_write, -) -from litellm.router_utils.auto_router_tuning_baseline import is_mutable_tuned_candidate, tuning_quota_violation from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.bedrock import AwsSessionTag from litellm.types.proxy.management_endpoints.model_management_endpoints import ( diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index 1f4ca442d20..41c4b5cfaf7 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -10,7 +10,7 @@ GET /router/fields - Get router settings field definitions without values (for U import inspect from collections.abc import Mapping from types import MappingProxyType -from typing import Any, Final, get_args +from typing import Final, get_args from fastapi import APIRouter, Depends from pydantic import Field @@ -32,7 +32,7 @@ router: Final = APIRouter() class RouterSettingsResponse(LiteLLMBaseModel): fields: list[RouterSettingsField] = Field(description="List of all configurable router settings with metadata") - current_values: dict[str, Any] = Field(description="Current values of router settings") + current_values: dict[str, object] = Field(description="Current values of router settings") routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option") routing_group_strategies: tuple[str, ...] = Field( description="Strategies supported when constructing a routing group" diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 1fa18e9b2b8..a256f417dfa 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -11,6 +11,7 @@ All /team management endpoints import asyncio import copy +import itertools import json import math import traceback @@ -139,7 +140,12 @@ from litellm.proxy.common_utils.user_api_key_cache import ( from litellm.proxy.db.db_span import db_span from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks from litellm.proxy.hooks.model_max_budget_limiter import ( + RoutedModels, build_model_max_budget_usage, + entries_outranked_by, + entry_names_model, + is_enforced_budget_entry, + proxy_routed_models, resolve_model_budget, ) from litellm.proxy.management.teams.authz import TEAM_OR_ORG_ADMIN, TeamRole, is_team_admin, team_access_denied @@ -1467,8 +1473,10 @@ def _check_team_model_budget_update_authority( data: UpdateTeamRequest, user_api_key_dict: UserAPIKeyAuth, existing_model_max_budget: Mapping[str, object] | None, + routed_models_source: Callable[[str | None], RoutedModels] = proxy_routed_models, ) -> None: - """Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap.""" + """Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap, or + add an entry that takes precedence over one.""" if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return if "model_max_budget" not in data.model_fields_set or not existing_model_max_budget: @@ -1487,29 +1495,97 @@ def _check_team_model_budget_update_authority( ) }, ) - for model_name, proposed in requested.items(): - governing = resolve_model_budget(model=model_name, model_max_budget=existing_model_max_budget) - if governing is None: - continue - cap = governing.budget_config - if cap.max_budget is None: - continue - if ( - proposed.max_budget is None - or proposed.max_budget > cap.max_budget - or proposed.budget_duration != cap.budget_duration - ): - raise HTTPException( - status_code=403, - detail={ - "error": ( - f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its " - f"budget_duration. Current max_budget={cap.max_budget} per {cap.budget_duration} " - f"(entry {governing.budget_model!r}), requested={proposed.max_budget} per " - f"{proposed.budget_duration}." - ) - }, + routed_models: Final = routed_models_source(data.team_id) + for entry, proposed in requested.items(): + existing = _existing_model_cap(existing_model_max_budget.get(entry)) + if existing is not None: + _check_existing_model_cap_kept(entry=entry, cap_entry=entry, proposed=proposed, existing=existing) + if is_enforced_budget_entry(existing) or not is_enforced_budget_entry(proposed): + continue + governing = resolve_model_budget(model=entry, model_max_budget=existing_model_max_budget, team_id=data.team_id) + if governing is not None and entry_names_model(entry=governing.budget_model, model=entry): + _check_existing_model_cap_kept( + entry=entry, cap_entry=governing.budget_model, proposed=proposed, existing=governing.budget_config ) + _check_new_entry_outranks_no_cap( + entry=entry, + existing_model_max_budget=existing_model_max_budget, + routed_models=routed_models, + team_id=data.team_id, + ) + + +def _check_existing_model_cap_kept(entry: str, cap_entry: str, proposed: BudgetConfig, existing: BudgetConfig) -> None: + if existing.max_budget is None or not _looser_model_cap(proposed=proposed, cap=existing): + return + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can raise a team's model_max_budget for {entry!r} or change its " + f"budget_duration. Current max_budget={existing.max_budget} per {existing.budget_duration} " + f"(entry {cap_entry!r}), requested={proposed.max_budget} per {proposed.budget_duration}." + ) + }, + ) + + +def _check_new_entry_outranks_no_cap( + entry: str, + existing_model_max_budget: Mapping[str, object], + routed_models: RoutedModels, + team_id: str | None, +) -> None: + """A new entry gets its own counter, so the models it takes from a capped entry stop drawing from that cap.""" + outranked: Final = next( + _capped_entries_outranked_by( + entry=entry, + existing_model_max_budget=existing_model_max_budget, + routed_models=routed_models, + team_id=team_id, + ), + None, + ) + if outranked is None: + return + outranked_name, cap = outranked + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can add a team model_max_budget entry {entry!r} that takes precedence over " + f"{outranked_name!r}, since the models they share would stop drawing from its max_budget=" + f"{cap.max_budget} per {cap.budget_duration}." + ) + }, + ) + + +def _capped_entries_outranked_by( + entry: str, + existing_model_max_budget: Mapping[str, object], + routed_models: RoutedModels, + team_id: str | None, +) -> Iterator[tuple[str, BudgetConfig]]: + governing: Final = resolve_model_budget(model=entry, model_max_budget=existing_model_max_budget, team_id=team_id) + outranked_names: Final = itertools.chain( + () + if governing is None or entry_names_model(entry=governing.budget_model, model=entry) + else (governing.budget_model,), + entries_outranked_by(entry=entry, model_max_budget=existing_model_max_budget, routed_models=routed_models), + ) + outranked_caps: Final = ((name, _existing_model_cap(existing_model_max_budget[name])) for name in outranked_names) + yield from ( + (name, cap) for name, cap in outranked_caps if name != entry and cap is not None and cap.max_budget is not None + ) + + +def _looser_model_cap(proposed: BudgetConfig, cap: BudgetConfig) -> bool: + return ( + proposed.max_budget is None + or (cap.max_budget is not None and proposed.max_budget > cap.max_budget) + or (is_enforced_budget_entry(cap) and proposed.budget_duration != cap.budget_duration) + ) def _should_auto_add_team_creator( @@ -1572,7 +1648,7 @@ async def new_team( - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) - - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} + - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only. @@ -2350,7 +2426,7 @@ async def update_team( - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) - - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} + - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only. diff --git a/litellm/proxy/management_helpers/auto_router_availability.py b/litellm/proxy/management_helpers/auto_router_availability.py index 9b511ce0817..ce7859aff2d 100644 --- a/litellm/proxy/management_helpers/auto_router_availability.py +++ b/litellm/proxy/management_helpers/auto_router_availability.py @@ -7,14 +7,14 @@ from typing import Final from pydantic import Json, TypeAdapter, ValidationError from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, capability_limit_violation, classify_strategy_router_model, count_capability_routers, gated_capabilities_of, ) -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( is_mutable_tuned_candidate, mutable_tuned_identities, tuning_quota_violation, diff --git a/litellm/proxy/management_helpers/auto_router_permissions.py b/litellm/proxy/management_helpers/auto_router_permissions.py index 995446db7a6..8597f0d3667 100644 --- a/litellm/proxy/management_helpers/auto_router_permissions.py +++ b/litellm/proxy/management_helpers/auto_router_permissions.py @@ -34,13 +34,16 @@ from litellm.repositories.prisma_protocols import DatabaseClient from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import TeamMembershipRepository from litellm.router import Router +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + classify_strategy_router_model, + strategy_router_dependencies, +) from litellm.router_strategy.complexity_router.config import ( ComplexityRouterConfigWrite, normalize_classifier_config_aliases, resolve_complexity_router_config_write, ) from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig -from litellm.router_utils.auto_router_model_naming import classify_strategy_router_model, strategy_router_dependencies from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import Deployment, updateDeployment diff --git a/litellm/proxy/management_helpers/team_roster_sync.py b/litellm/proxy/management_helpers/team_roster_sync.py index dc0ba2e47d4..08be455a55d 100644 --- a/litellm/proxy/management_helpers/team_roster_sync.py +++ b/litellm/proxy/management_helpers/team_roster_sync.py @@ -40,6 +40,7 @@ from litellm.proxy.management_helpers.access_group_team_sync import ( ) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.chunked_in import delete_many_in, find_many_in, update_many_in +from litellm.repositories.team_repository import TeamRepository if TYPE_CHECKING: from prisma import Prisma @@ -144,7 +145,7 @@ def _changes_nothing(plan: RosterPlan) -> bool: async def _unchanged_roster(prisma_client: PrismaClient, team_id: str) -> RosterSync | TeamGone: - team_row: Final = await prisma_client.writer_db.litellm_teamtable.find_unique(where={"team_id": team_id}) + team_row: Final = await TeamRepository(prisma_client, use_writer=True).table.find_unique(where={"team_id": team_id}) if team_row is None: return TeamGone(team_id=team_id) team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump()) diff --git a/litellm/proxy/mcp_server/discoverable_endpoints.py b/litellm/proxy/mcp_server/discoverable_endpoints.py index eb12b15cdae..07a267d6caa 100644 --- a/litellm/proxy/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/mcp_server/discoverable_endpoints.py @@ -300,12 +300,13 @@ def decode_state_hash(encrypted_state: str) -> dict: if decrypted_json is None: raise ValueError("Failed to decrypt state parameter") - state_data: Final = json.loads(decrypted_json) + state_data: Final = _DECODED_STATE.validate_python(json.loads(decrypted_json)) return state_data _MANAGED_AUTH_CODE_PREFIX: Final = "llm_scope_" _DECODED_JSON: Final = TypeAdapter(object) +_DECODED_STATE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) class _ManagedAuthorizationCode(LiteLLMBaseModel): diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index d2ceefffe47..5d63c3db222 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -6,9 +6,12 @@ ###################################################################### import asyncio +import functools +import itertools import traceback -from collections.abc import Mapping, Sequence -from typing import Any, BinaryIO, Final, TypedDict, cast, get_args +from collections.abc import Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from typing import Any, BinaryIO, Final, NoReturn, TypedDict, cast, get_args import httpx from fastapi import ( @@ -28,12 +31,13 @@ from typing_extensions import ReadOnly import litellm from litellm import CreateFileRequest, get_secret_str from litellm._logging import verbose_proxy_logger +from litellm.constants import ANTHROPIC_DEFAULT_FILE_LIST_LIMIT from litellm.litellm_core_utils.cloud_storage_security import ( is_managed_cloud_storage_uri, ) from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME -from litellm.llms.base_llm.files.transformation import BaseFileEndpoints +from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope from litellm.llms.base_llm.managed_resources.isolation import build_list_page from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -112,6 +116,9 @@ from litellm.types.llms.openai import ( OpenAIFileObject, OpenAIFilesPurpose, ) +from litellm.types.proxy.auth.user_api_key_auth import user_api_key_has_admin_view +from litellm.types.router import Deployment +from litellm.types.utils import LlmProviders router: Final = APIRouter() @@ -1673,6 +1680,185 @@ def _as_file_list_page(response: object) -> object: return FileListPage.model_validate(page) +async def _list_owned_managed_files( + managed_files_obj: BaseFileEndpoints, + user_api_key_dict: UserAPIKeyAuth, + purpose: str | None, + limit: int | None, + after: str | None, + before: str | None, + model_scope: FileListModelScope | None, +) -> FileListPage: + return await managed_files_obj.afile_list( + purpose=purpose, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + user_api_key_dict=user_api_key_dict, + limit=limit, + after=after, + before=before, + model_scope=model_scope, + ) + + +@dataclass(frozen=True, slots=True) +class _UnscopedFileList: + status_code: int + detail: str + + +@dataclass(frozen=True, slots=True) +class _ListedFiles: + model_scope: FileListModelScope | None + custom_llm_provider: str | None + + +def _raise_unscoped_file_list(failure: _UnscopedFileList) -> NoReturn: + raise HTTPException(status_code=failure.status_code, detail=failure.detail) + + +def _single_target_model_name(target_model_names: str) -> str | _UnscopedFileList: + target_model_names_list: Final = target_model_names.split(",") + if len(target_model_names_list) != 1: + return _UnscopedFileList( + status_code=400, + detail="target_model_names on list files must be a list of one model name. Example: ['gpt-4o']", + ) + return target_model_names_list[0] + + +async def _requested_custom_llm_provider(request: Request, provider: str | None) -> str | None: + return ( + provider + or get_custom_llm_provider_from_request_headers(request=request) + or get_custom_llm_provider_from_request_query(request=request) + or await get_custom_llm_provider_from_request_body(request=request) + ) + + +def _default_file_list_limit(custom_llm_provider: str | None) -> int | None: + return ANTHROPIC_DEFAULT_FILE_LIST_LIMIT if custom_llm_provider == LlmProviders.ANTHROPIC.value else None + + +def _deployment_model_mapping_keys(deployment: Deployment) -> frozenset[str]: + return frozenset(key for key in (deployment.model_name, deployment.model_info.id) if key is not None) + + +def _model_mapping_keys(deployments: Iterable[Deployment]) -> frozenset[str]: + return frozenset(itertools.chain.from_iterable(map(_deployment_model_mapping_keys, deployments))) + + +def _deployments_serving_model(llm_router: Router, model: str) -> tuple[Deployment, ...]: + deployment: Final = llm_router.get_deployment(model_id=model) if llm_router.has_model_id(model) else None + if deployment is not None: + return (deployment,) + return tuple(Deployment.model_validate(listed) for listed in llm_router.get_model_list(model_name=model) or ()) + + +def _resolved_deployment_provider(deployment: Deployment) -> str | None: + try: + return litellm.get_llm_provider( + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + )[1] + except litellm.BadRequestError: + return None + + +def _deployments_on_provider(llm_router: Router | None, custom_llm_provider: str) -> Iterator[Deployment]: + if llm_router is None: + return iter(()) + return ( + deployment + for deployment_id in llm_router.get_model_ids() + if (deployment := llm_router.get_deployment(model_id=deployment_id)) is not None + and _resolved_deployment_provider(deployment) == custom_llm_provider + ) + + +def _routed_model_name_on_provider( + llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str +) -> bool: + if model_name in llm_router.model_names: + return model_name in provider_model_names + return llm_router.pattern_router.routes_to_provider(model_name, custom_llm_provider) + + +def _stored_model_name_routes_to_provider( + llm_router: Router, custom_llm_provider: str, provider_model_names: frozenset[str], model_name: str +) -> bool: + routing_group: Final = llm_router.get_routing_group(model_name) + routed_model_names: Final = ( + tuple(routing_group.models) if routing_group is not None else (llm_router.routable_model_group(model_name),) + ) + return any( + _routed_model_name_on_provider(llm_router, custom_llm_provider, provider_model_names, routed_model_name) + for routed_model_name in routed_model_names + ) + + +def _files_on_provider(llm_router: Router | None, custom_llm_provider: str) -> _ListedFiles | _UnscopedFileList: + from litellm.proxy.pass_through_endpoints.managed_id_rewriter import passthrough_file_model_id + + if custom_llm_provider not in litellm.provider_list: + return _UnscopedFileList(status_code=400, detail=f"Unknown provider '{custom_llm_provider}' for listing files.") + passthrough_model_id: Final = passthrough_file_model_id(custom_llm_provider) + deployments: Final = tuple(_deployments_on_provider(llm_router, custom_llm_provider)) + model_mapping_keys: Final = _model_mapping_keys(deployments) | frozenset( + (passthrough_model_id,) if passthrough_model_id is not None else () + ) + if llm_router is None: + return _ListedFiles(FileListModelScope(model_mapping_keys), custom_llm_provider) + return _ListedFiles( + FileListModelScope( + model_mapping_keys, + routes_stored_model_name=functools.partial( + _stored_model_name_routes_to_provider, + llm_router, + custom_llm_provider, + frozenset(deployment.model_name for deployment in deployments), + ), + ), + custom_llm_provider, + ) + + +def _files_serving_model(llm_router: Router, model: str) -> _ListedFiles: + deployments: Final = _deployments_serving_model(llm_router, model) + providers: Final = frozenset(map(_resolved_deployment_provider, deployments)) + return _ListedFiles( + FileListModelScope(_model_mapping_keys(deployments)), next(iter(providers)) if len(providers) == 1 else None + ) + + +async def _listed_files( + llm_router: Router | None, + user_api_key_dict: UserAPIKeyAuth, + custom_llm_provider: str | None, + target_model_names: str | None, + routed_model: str | None, +) -> _ListedFiles | _UnscopedFileList: + if routed_model is not None and llm_router is not None: + return _files_serving_model(llm_router, routed_model) + if target_model_names: + target_model_name: Final = _single_target_model_name(target_model_names) + if isinstance(target_model_name, _UnscopedFileList): + return target_model_name + if llm_router is None: + return _UnscopedFileList( + status_code=500, detail="LLM Router not initialized. Ensure models added to proxy." + ) + await get_authorized_credentials_for_model( + llm_router=llm_router, + model_id=target_model_name, + user_api_key_dict=user_api_key_dict, + operation_context="file list", + ) + return _files_serving_model(llm_router, target_model_name) + if custom_llm_provider is None: + return _ListedFiles(None, None) + return _files_on_provider(llm_router, custom_llm_provider) + + @router.get( "/{provider}/v1/files", dependencies=[Depends(user_api_key_auth)], @@ -1697,6 +1883,8 @@ async def list_files( purpose: str | None = None, limit: int | None = None, after: str | None = None, + after_id: str | None = None, + before_id: str | None = None, ): """ Returns information about a specific file. that can be used across - Assistants API, Batch API @@ -1739,6 +1927,10 @@ async def list_files( ) response: Any | None = None + owned_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files") + list_owned_files_only: Final = litellm.require_managed_files is True and not user_api_key_has_admin_view( + user_api_key_dict + ) # Check for model-based credential routing (no file_id encoding check for list) should_route, model_used, _, credentials = await handle_model_based_routing( @@ -1750,7 +1942,36 @@ async def list_files( check_file_id_encoding=False, ) - if should_route and credentials is not None: + if list_owned_files_only: + if after and after_id and after != after_id: + raise HTTPException( + status_code=400, + detail="Pass the list cursor as either 'after' or 'after_id', not both with different ids.", + ) + if not isinstance(owned_files_hook, BaseFileEndpoints): + raise HTTPException( + status_code=500, + detail="Managed file listing is unavailable, so require_managed_files cannot scope the list to the caller.", + ) + listed_files: Final = await _listed_files( + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=await _requested_custom_llm_provider(request=request, provider=provider), + target_model_names=target_model_names, + routed_model=model_used if should_route else None, + ) + if isinstance(listed_files, _UnscopedFileList): + _raise_unscoped_file_list(listed_files) + response = await _list_owned_managed_files( # rebind-ok: each dispatch branch binds response + managed_files_obj=owned_files_hook, + user_api_key_dict=user_api_key_dict, + purpose=purpose, + limit=limit if limit is not None else _default_file_list_limit(listed_files.custom_llm_provider), + after=after or after_id, + before=before_id, + model_scope=listed_files.model_scope, + ) + elif should_route and credentials is not None: # Use model-based routing with credentials from config prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True) response = await litellm.afile_list( @@ -1762,12 +1983,9 @@ async def list_files( verbose_proxy_logger.debug("Listed files using model: %s", model_used) elif target_model_names and isinstance(target_model_names, str): - target_model_names_list: Final = target_model_names.split(",") - if len(target_model_names_list) != 1: - raise HTTPException( - status_code=400, - detail="target_model_names on list files must be a list of one model name. Example: ['gpt-4o']", - ) + target_model_name: Final = _single_target_model_name(target_model_names) + if isinstance(target_model_name, _UnscopedFileList): + _raise_unscoped_file_list(target_model_name) if llm_router is None: raise HTTPException( status_code=500, @@ -1775,7 +1993,7 @@ async def list_files( ) credentials = await get_authorized_credentials_for_model( llm_router=llm_router, - model_id=target_model_names_list[0], + model_id=target_model_name, user_api_key_dict=user_api_key_dict, operation_context="file list", ) @@ -1786,20 +2004,17 @@ async def list_files( **data, ) else: - custom_llm_provider: Final = ( - provider - or get_custom_llm_provider_from_request_headers(request=request) - or get_custom_llm_provider_from_request_query(request=request) - or await get_custom_llm_provider_from_request_body(request=request) - ) + custom_llm_provider: Final = await _requested_custom_llm_provider(request=request, provider=provider) managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") if custom_llm_provider is None and isinstance(managed_files_obj, BaseFileEndpoints): - response = await managed_files_obj.afile_list( - purpose=purpose, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + response = await _list_owned_managed_files( # rebind-ok: each dispatch branch binds response + managed_files_obj=managed_files_obj, user_api_key_dict=user_api_key_dict, + purpose=purpose, limit=limit, after=after, + before=None, + model_scope=None, ) else: resolved_custom_llm_provider: Final = custom_llm_provider or "openai" diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 7c7eb51bc5d..5415314f29b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -29,6 +29,7 @@ from litellm.llms.anthropic.messages.stream_assembly import ( split_sse_chunk_into_events, ) from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -61,7 +62,7 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_i def _event_payload(event: str) -> Mapping[str, object] | None: data_lines: Final = tuple( - line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:") + line.strip()[len("data:") :].strip() for line in split_sse_lines(event) if line.strip().startswith("data:") ) if not data_lines or data_lines[0] in ("", "[DONE]"): return None @@ -362,7 +363,7 @@ class AnthropicPassthroughLoggingHandler: ) -> str | None: for raw in all_chunks: text = raw.decode("utf-8") if isinstance(raw, bytes) else raw - for line in text.splitlines(): + for line in split_sse_lines(text): if not line.startswith("data:"): continue try: @@ -402,7 +403,7 @@ class AnthropicPassthroughLoggingHandler: """ for raw in reversed(all_chunks): text = raw.decode("utf-8") if isinstance(raw, bytes) else raw - for line in reversed(text.splitlines()): + for line in reversed(split_sse_lines(text)): if not line.startswith("data:"): continue try: diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py index 2afc1d17bc2..0c4a3f64852 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py @@ -21,6 +21,7 @@ from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.tinyfish import ( @@ -81,7 +82,7 @@ def sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> bool: def run_id_from_sse_frames(frames: bytes) -> str | None: - return _run_id_from_sse_chunks(frames.decode("utf-8", errors="replace").splitlines()) + return _run_id_from_sse_chunks(split_sse_lines(frames.decode("utf-8", errors="replace"))) def resolve_tinyfish_agent_api_base() -> str: diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index a5a04fb68d2..cae8ccc93ee 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -51,6 +51,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( can_access_resource, resolve_resource_owner_id, ) +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames @@ -242,6 +243,12 @@ def _passthrough_sentinel_model_id(provider: str) -> str: return f"_passthrough_{provider}" +def passthrough_file_model_id(provider: str) -> str | None: + if (provider, "/v1/files") not in _LIST_ROUTE_TABLE: + return None + return _passthrough_sentinel_model_id(provider) + + # Key under which the provider marker is stored in a file row's model_mappings. # Its value lands in flat_model_file_ids (built from model_mappings.values()), # giving the file table a DB-queryable provider scope it otherwise lacks. @@ -824,7 +831,7 @@ _SSE_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue]) def _first_streamed_response(frames: bytes) -> tuple[str, Mapping[str, JsonValue]] | None: - for line in frames.decode("utf-8", errors="replace").splitlines(): + for line in split_sse_lines(frames.decode("utf-8", errors="replace")): if not line.startswith(_SSE_DATA_PREFIX): continue try: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 650dcae3239..e017eee3d27 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -30,6 +30,7 @@ from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Route from starlette.websockets import WebSocketState +from typing_extensions import ReadOnly from websockets.asyncio.client import connect from websockets.exceptions import ( ConnectionClosedError, @@ -173,8 +174,17 @@ def _is_mergeable_client_metadata(key: str, value: object) -> bool: return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping)) +class _RegisteredPassThroughRoute(TypedDict): + endpoint_id: ReadOnly[str] + path: ReadOnly[str | None] + type: ReadOnly[str] + methods: ReadOnly[str | bool | list[str] | tuple[str, ...] | Mapping[str, object]] + auth: ReadOnly[bool] + passthrough_params: ReadOnly[Mapping[str, object]] + + # Global registry to track registered pass-through routes and prevent memory leaks -_registered_pass_through_routes: Final[dict[str, dict[str, str | bool | list[str] | Mapping[str, object]]]] = {} +_registered_pass_through_routes: Final[dict[str, _RegisteredPassThroughRoute]] = {} def get_response_body(response: httpx.Response) -> dict | None: @@ -2239,7 +2249,7 @@ def create_pass_through_route( status_code=status.HTTP_405_METHOD_NOT_ALLOWED, detail=f"Method {request.method} is not allowed for pass-through endpoint {path}.", ) - target_params: Final = { + target_params: Final[dict[str, object]] = { "target": target, "custom_headers": custom_headers, "forward_headers": _forward_headers, @@ -2271,7 +2281,9 @@ def create_pass_through_route( # Ensure custom_headers is a dict. Botocore returns a HeadersDict # for SigV4-prepared requests, which is a Mapping but not a dict. - headers_dict: Final = dict(param_custom_headers) if isinstance(param_custom_headers, Mapping) else {} + headers_dict: Final[dict[str, object]] = ( + {**param_custom_headers} if isinstance(param_custom_headers, Mapping) else {} + ) # Ensure query_params and custom_body are dicts or None final_query_params: Final = query_params_data if isinstance(query_params_data, dict) else {} @@ -3104,9 +3116,7 @@ class SafeRouteAdder: return f"{path}/{{subpath:path}}" return None - def registered_route_methods( - route_info: Mapping[str, str | bool | Sequence[str] | Mapping[str, object]], - ) -> tuple[tuple[str, str], ...]: + def registered_route_methods(route_info: _RegisteredPassThroughRoute) -> tuple[tuple[str, str], ...]: starlette_path: Final = registered_route_path(route_info) methods: Final = route_info.get("methods") if starlette_path is None or not isinstance(methods, (list, tuple)): @@ -3462,7 +3472,7 @@ class InitPassThroughEndpointHelpers: return (comparison_route, *forwarded) @staticmethod - def get_registered_pass_through_route(route: str, method: str | None = None) -> dict[str, Any] | None: + def get_registered_pass_through_route(route: str, method: str | None = None) -> _RegisteredPassThroughRoute | None: """Get passthrough params for a given route and optionally filter by HTTP method""" comparison_route: Final = InitPassThroughEndpointHelpers._route_for_registry_lookup(route) for key in _registered_pass_through_routes: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 64fb2791913..b263412ce0f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -168,12 +168,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields -from litellm.router_utils.access_windows import access_windows_config_error -from litellm.router_utils.add_retry_fallback_headers import ( - get_fallback_errors_from_headers, - get_hidden_params_dict, -) -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, capability_limit_violation, @@ -181,12 +176,17 @@ from litellm.router_utils.auto_router_model_naming import ( count_capability_routers, validate_complexity_router_config_placement, ) -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( TUNING_BASELINE_PARAM_NAME, mutable_tuned_identities, snapshot_tuning_baselines, tuning_limit_violation, ) +from litellm.router_utils.access_windows import access_windows_config_error +from litellm.router_utils.add_retry_fallback_headers import ( + get_fallback_errors_from_headers, + get_hidden_params_dict, +) from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.router_utils.routing_groups import parse_routing_groups from litellm.types.caching import RedisPipelineIncrementOperation @@ -666,7 +666,11 @@ from litellm.proxy.list_api.common import ( problem_response, request_validation_problem, ) -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_litellm_data_to_request +from litellm.proxy.litellm_pre_call_utils import ( + SCHEDULER_PRIORITY_HEADER, + LiteLLMProxyRequestSetup, + add_litellm_data_to_request, +) from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( rust_control_plane_router, ) @@ -18481,6 +18485,11 @@ async def async_queue_request( raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}) router: Final = llm_router + data.update( + router.priority_header_update( + model=data["model"], kwargs=data, header=request.headers.get(SCHEDULER_PRIORITY_HEADER) + ) + ) queued_priority, _ = router.split_request_priority(model=data["model"], kwargs=data) if "stream" in data and data["stream"] is True: # use generate_responses to stream responses diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 0dd96c856f8..f2d88702470 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -4145,6 +4145,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "ZeroGPU", + "provider_display_name": "ZeroGPU", + "litellm_provider": "zerogpu", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.zerogpu.ai/v1", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "zerogpu/gpt-oss-120b" + }, { "provider": "CURSOR", "provider_display_name": "Cursor", diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 5f6a7432afb..1eb76f8f558 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -18,6 +18,7 @@ from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler @@ -81,7 +82,7 @@ class _StreamsBody(Protocol): def _sse_frame_data(frame: str) -> str | None: - return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None) + return next((line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")), None) async def _never_receive() -> Message: diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index cb226929167..db1951ec9e2 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -256,7 +256,6 @@ REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = Ma "avideo_create_character": ("name", "video"), "acreate_container": ("name",), "aupload_container_file": ("file",), - "acreate_agent": ("name",), "acreate_interaction": ("input",), "acreate_eval": ("data_source_config", "testing_criteria"), "acreate_run": ("data_source",), @@ -264,7 +263,7 @@ REQUIRED_PRESENT_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, ...]]] = Ma ) REQUIRED_ONE_OF_BODY_PARAMS_BY_ROUTE: Final[Mapping[str, tuple[str, str]]] = MappingProxyType( - {"acreate_interaction": ("model", "agent")} + {"acreate_interaction": ("model", "agent"), "acreate_agent": ("id", "name")} ) JSON_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 83a0e6f9e05..2965174114b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) @@index([flat_model_file_ids], type: Gin) @@index([team_id, created_at(sort: Desc)]) + @@index([created_by, created_at(sort: Desc)]) } model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the @@ -2060,3 +2061,73 @@ model LiteLLM_TelemetryReport { @@index([window_end]) @@index([created_at]) } + +enum AccountingEventKind { + inbound_usage + outbound_usage + cost +} + +enum AccountingDeliveryState { + pending + applied + parked +} + +enum AccountingDeliveryTarget { + LiteLLM_VerificationToken + LiteLLM_UserTable + LiteLLM_TeamTable + LiteLLM_TeamMembership + LiteLLM_OrganizationTable + LiteLLM_OrganizationMembership + LiteLLM_EndUserTable + LiteLLM_AgentsTable + LiteLLM_TagTable + LiteLLM_ProjectTable + LiteLLM_ModelAccessGroupBudgetTable + LiteLLM_BudgetTable + LiteLLM_BudgetWindowSpend + LiteLLM_DailyUserSpend + LiteLLM_DailyTeamSpend + LiteLLM_DailyOrganizationSpend + LiteLLM_DailyEndUserSpend + LiteLLM_DailyAgentSpend + LiteLLM_DailyTagSpend + LiteLLM_DailyModelUsage + LiteLLM_DailyToolSpend + LiteLLM_AutoRouterDailySpend +} + +// Durable accounting ledger. One row per usage or cost event; payload holds the +// private event envelope versioned by schema_version. Usage rows retain request-scoped +// pricing context; cost rows include resolved rates and the finalized pricing snapshot. +model LiteLLM_AccountingEvent { + event_id String @id + kind AccountingEventKind + source_event_id String? + schema_version Int + payload_hash String + occurred_at DateTime + recorded_at DateTime @default(now()) + payload Json + + @@index([recorded_at]) + @@index([kind, source_event_id]) +} + +// One row per (event, target table, version): pending work for the projector and, +// once applied in the same transaction as the target write, the receipt that stops a retry. +model LiteLLM_AccountingDelivery { + event_id String + target AccountingDeliveryTarget + version Int + state AccountingDeliveryState @default(pending) + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + + @@id([event_id, target, version]) + @@index([target, state, next_attempt_at]) +} diff --git a/litellm/proxy/spend_tracking/key_metadata_recovery.py b/litellm/proxy/spend_tracking/key_metadata_recovery.py index 5442c0f3f81..8fd4fbfcb8b 100644 --- a/litellm/proxy/spend_tracking/key_metadata_recovery.py +++ b/litellm/proxy/spend_tracking/key_metadata_recovery.py @@ -6,7 +6,7 @@ from datetime import datetime, timedelta from types import MappingProxyType from typing import Final, Literal, TypeVar -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -155,6 +155,7 @@ class _DailyUserSpendOwnerRow(LiteLLMBaseModel): last_owner: str | None = None +_RAW_ROWS: Final = TypeAdapter(Sequence[dict[str, object]], config=ConfigDict(strict=True)) _TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...]) _SPEND_LOG_DIGEST_ROWS: Final = TypeAdapter(tuple[_SpendLogDigestRow, ...]) _DAILY_USER_SPEND_OWNER_ROWS: Final = TypeAdapter(tuple[_DailyUserSpendOwnerRow, ...]) @@ -194,7 +195,7 @@ async def _rows_within_the_statement_timeout( await transaction.execute_raw(_SPEND_LOG_STATEMENT_TIMEOUT_SQL) for setting in planner_settings: await transaction.execute_raw(setting) - return await transaction.query_raw(sql, *params) + return _RAW_ROWS.validate_python(await transaction.query_raw(sql, *params)) async def _reverse_hash_key_metadata( diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 10c603fe56f..4ca23e4e89e 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -19,6 +19,7 @@ from litellm.repositories.daily_activity_sql import ( build_cache_leakage_keys_sql, build_entity_rollup_sql, build_export_sql, + build_guardrail_spend_sql, build_key_page_sql, build_key_search_sql, build_model_top_keys_sql, @@ -36,6 +37,7 @@ from litellm.types.repositories.daily_activity import ( ExportRow, ExportType, GroupingSetsRow, + GuardrailSpendRow, KeyMetadataRow, KeyPage, KeySpendRow, @@ -100,6 +102,7 @@ class DailyActivityDatabase(Protocol): _GROUPING_ADAPTER: Final = TypeAdapter(tuple[GroupingSetsRow, ...]) _ENTITY_ADAPTER: Final = TypeAdapter(tuple[EntityRollupRow, ...]) +_GUARDRAIL_SPEND_ADAPTER: Final = TypeAdapter(tuple[GuardrailSpendRow, ...]) _KEY_SPEND_ADAPTER: Final = TypeAdapter(tuple[KeySpendRow, ...]) _KEY_PAGE_TOTAL_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(int) _USER_SPEND_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...]) @@ -191,6 +194,9 @@ class DailyActivityRepository: distinct_api_keys=distinct_api_keys, ) + async def guardrail_spend(self, scope: DailyActivityScope) -> tuple[GuardrailSpendRow, ...]: + return _GUARDRAIL_SPEND_ADAPTER.validate_python(await self._query(build_guardrail_spend_sql(scope))) + async def search_keys(self, scope: DailyActivityScope, *, search: str, limit: int) -> tuple[str, ...]: if not 1 <= limit <= constants.USAGE_KEY_SEARCH_MAX: raise ValueError(f"limit must be between 1 and {constants.USAGE_KEY_SEARCH_MAX}") diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index 3d27e7d937e..7ec9fa67a67 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -261,6 +261,25 @@ def build_entity_rollup_sql(scope: DailyActivityScope, *, api_key_limit: int) -> return SqlQuery(sql=sql, params=(*where_params, PTU_SENTINEL_API_KEY, api_key_limit)) +def build_guardrail_spend_sql(scope: DailyActivityScope) -> SqlQuery: + pg_table: Final = PRISMA_TO_PG_TABLE[scope.table] + where_clause, where_params = build_where_clause(scope) + sql: Final = f""" + SELECT g.date, g.guardrail_id, COALESCE(t.guardrail_name, g.guardrail_id) AS guardrail_name, g.api_key, + COALESCE(SUM(g.cost), 0)::float AS spend, + COALESCE(SUM(CASE WHEN g.cost IS NULL THEN g.units ELSE g.untracked_units END), 0)::bigint AS untracked_units + FROM "LiteLLM_DailyGuardrailUsageUnits" g + LEFT JOIN "LiteLLM_GuardrailsTable" t ON t.guardrail_id = g.guardrail_id + WHERE g.date >= $1 AND g.date <= $2 + AND EXISTS ( + SELECT 1 FROM "{pg_table}" + WHERE {where_clause} AND "{pg_table}".date = g.date AND "{pg_table}".api_key = g.api_key + ) + GROUP BY g.date, g.guardrail_id, t.guardrail_name, g.api_key + """ + return SqlQuery(sql=sql, params=where_params) + + def _key_spend_select() -> str: return """ COALESCE(SUM(spend), 0)::float AS spend, diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 219944de521..982158040c4 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -165,6 +165,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._upstream_exhausted: bool = False self._response_id_primed: bool = False self._pending_tool_events: list[BaseLiteLLMOpenAIResponseObject] = [] + self._tool_calls_held_behind_text: tuple[object, ...] = () self._tool_output_index_by_call_id: dict[str, int] = {} self._tool_args_by_call_id: dict[str, str] = {} self._tool_item_id_by_call_id: dict[str, str] = {} # mutable-ok: filled per call id as tool call events stream @@ -388,6 +389,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return "" return self._tool_call_id_by_index.get(index, "") + def _hold_tool_calls_behind_text(self, tool_calls: object) -> None: + if isinstance(tool_calls, list): + self._tool_calls_held_behind_text = (*self._tool_calls_held_behind_text, *tool_calls) + + def _release_held_tool_calls(self) -> tuple[object, ...]: + held: Final = self._tool_calls_held_behind_text + self._tool_calls_held_behind_text = () + return held + def _queue_tool_call_delta_events(self, tool_calls: object) -> None: """ Convert chat-completions streaming `tool_calls` deltas into Responses API streaming events. @@ -1365,6 +1375,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if delta_content: if not self.sent_message_item_added_event: self._queue_message_item_added_events() + self._hold_tool_calls_behind_text(getattr(chunk.choices[0].delta, "tool_calls", None)) self._sequence_number += 1 text_delta_event: Final = OutputTextDeltaEvent( type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, @@ -1379,7 +1390,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): # Priority 3: Handle tool call deltas (if any) -> queue events and emit them # For each tool call delta, we emit events one at a time to match OpenAI's streaming behavior if chunk.choices and hasattr(chunk.choices[0].delta, "tool_calls") and chunk.choices[0].delta.tool_calls: - self._queue_tool_call_delta_events(chunk.choices[0].delta.tool_calls) + self._queue_tool_call_delta_events([*self._release_held_tool_calls(), *chunk.choices[0].delta.tool_calls]) # Return one pending tool event at a time if self._pending_tool_events: return self._pending_tool_events.pop(0) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c1f422c7d7b..4e418a91156 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -41,7 +41,11 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base -from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation +from litellm.litellm_core_utils.internal_call_metadata import ( + is_background_response, + requests_background, + without_budget_reservation, +) from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base @@ -65,7 +69,7 @@ from litellm.types.llms.openai import ( ResponsesAPIStreamingResponse, ) from litellm.types.responses.main import GenericResponseOutputItem, OutputText -from litellm.types.utils import CallTypes +from litellm.types.utils import CallTypes, ModelResponse from litellm.utils import async_post_call_success_deployment_hook if TYPE_CHECKING: @@ -288,7 +292,7 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: return _status_code_for_error_fields(error_type, error_code), message -def _map_stream_error_to_exception( +def map_stream_error_to_exception( error_obj: object, model: str, custom_llm_provider: str, @@ -578,6 +582,7 @@ class BaseResponsesAPIStreamingIterator: self.request_data.get("input"), self.request_data, self._generated_content + self._generated_tool_arguments, + use_default_image_token_count=False, ) if _estimate_wanted else None @@ -731,7 +736,7 @@ class BaseResponsesAPIStreamingIterator: self._earlier_rounds_spend = earlier def _map_error_event_exception(self, error_obj: object) -> Exception: - return _map_stream_error_to_exception( + return map_stream_error_to_exception( error_obj, self.model or "", self.custom_llm_provider or "", @@ -981,6 +986,39 @@ class BaseResponsesAPIStreamingIterator: except Exception: pass + def _record_partial_usage_for_failure(self) -> None: + generated_text: Final = self._generated_content + self._generated_tool_arguments + if ( + not generated_text + or requests_background(self.request_data) + or is_background_response(self._created_response) + or self._completed_response_logged + or self.logging_obj.model_call_details.get("combined_usage_object") is not None + ): + return + estimate: Final = _estimate_usage_safely( + self.model or "", + self.request_data.get("input"), + self.request_data, + generated_text, + use_default_image_token_count=True, + ) + if estimate is None: + return + usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the success path above uses + estimate + ) + served_tier: Final[object] = getattr(self._created_response, "service_tier", None) + assembled: Final = ModelResponse( + model=self.model, usage=usage, service_tier=served_tier if isinstance(served_tier, str) else None + ) + if self._earlier_rounds_spend is None: + self.logging_obj.record_assembled_response_for_failure(assembled) + return + self.logging_obj.record_failure_after_earlier_rounds( + self._earlier_rounds_spend, usage, self.logging_obj.response_cost_calculator(result=assembled) or 0.0 + ) + def _handle_failure(self, exception: Exception): """ Trigger failure handlers before bubbling the exception. @@ -990,6 +1028,7 @@ class BaseResponsesAPIStreamingIterator: if self._failure_handled: return self._failure_handled = True + self._record_partial_usage_for_failure() traceback_exception: Final = traceback.format_exc() end_time: Final = datetime.now() @@ -1060,15 +1099,20 @@ class BaseResponsesAPIStreamingIterator: the same shape the success handler assembles for a natural end of stream. Its response is the response.created payload (it carries the served service_tier) plus the text and tool arguments relayed so far, with counted - usage since the provider never sent any. None when nothing was relayed yet - or the stream already reached its terminal event and billed itself. + usage since the provider never sent any. None when nothing was relayed yet, + the stream already reached its terminal event and billed itself, or its failure + log already billed the relayed output. """ created: Final = self._created_response generated: Final = self._generated_content + self._generated_tool_arguments - if self.completed_response is not None or created is None or not generated: + if created is None or not generated or self._billed_without_disconnect(): return None usage: Final = _estimate_usage_safely( - self.model or "", self.request_data.get("input"), self.request_data, generated + self.model or "", + self.request_data.get("input"), + self.request_data, + generated, + use_default_image_token_count=False, ) if usage is None: return None @@ -1082,6 +1126,9 @@ class BaseResponsesAPIStreamingIterator: stamp_responses_usage_cost(partial, self.logging_obj) return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial) + def _billed_without_disconnect(self) -> bool: + return self.completed_response is not None or self._failure_handled + def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: self._yielded_first_chunk = True if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: @@ -1218,6 +1265,9 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): raise StopAsyncIteration from e self._raise_for_transport_error(e) except httpx.TimeoutException as e: + if self.completed_response is not None: + self.finished = True + raise StopAsyncIteration from e self._raise_for_transport_timeout(e) except httpx.HTTPError as e: # Handle HTTP errors @@ -1303,6 +1353,9 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): raise StopIteration from e self._raise_for_transport_error(e) except httpx.TimeoutException as e: + if self.completed_response is not None: + self.finished = True + raise StopIteration from e self._raise_for_transport_timeout(e) except httpx.HTTPError as e: # Handle HTTP errors @@ -1727,6 +1780,8 @@ def _estimate_usage_from_text( request_input: object, responses_api_request: Mapping[str, object], generated_text: str, + *, + use_default_image_token_count: bool = False, ) -> ResponseAPIUsage: messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union @@ -1734,7 +1789,7 @@ def _estimate_usage_from_text( only_since_compaction=True, ) input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped - model=model, messages=messages + model=model, messages=messages, use_default_image_token_count=use_default_image_token_count ) output_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped model=model, text=generated_text, count_response_tokens=True @@ -1751,6 +1806,8 @@ def _estimate_usage_safely( request_input: object, responses_api_request: Mapping[str, object], generated_text: str, + *, + use_default_image_token_count: bool, ) -> ResponseAPIUsage | None: try: return _estimate_usage_from_text( @@ -1758,6 +1815,7 @@ def _estimate_usage_safely( request_input=request_input, responses_api_request=responses_api_request, generated_text=generated_text, + use_default_image_token_count=use_default_image_token_count, ) except Exception as e: verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e) @@ -2276,7 +2334,7 @@ class _ResponsesWsTurn: ) if failed_event is None: return None - return _map_stream_error_to_exception( + return map_stream_error_to_exception( _ws_event_error(failed_event), model, custom_llm_provider, logging_obj=connection_logging_obj ) diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 02d07932533..cf993c69814 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -1,8 +1,10 @@ import base64 +import os import re from collections.abc import Callable, Collection, Iterable, Mapping, Sequence from functools import reduce -from typing import Any, Final, Optional, TypeVar, Union, cast, get_type_hints, overload +from types import MappingProxyType +from typing import Any, Final, Optional, TypeVar, Union, cast, get_args, get_type_hints, overload from pydantic import BaseModel, TypeAdapter from typing_extensions import TypedDict as ExtraItemsTypedDict @@ -14,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value, is_nested_path from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig, ReasoningSummaryLevel +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import ( AllMessageValues, InputTokensDetails, @@ -29,6 +32,7 @@ from litellm.types.responses.main import DecodedResponseId from litellm.types.utils import ( CacheCreationTokenDetails, CompletionTokensDetailsWrapper, + ModelInfo, PromptTokensDetailsWrapper, ProviderSpecificHeader, SpecialEnums, @@ -1382,3 +1386,98 @@ def auto_reasoning_summary(model: str, custom_llm_provider: object) -> Reasoning provider, model.removeprefix(f"{provider}/") ) return "detailed" if responses_config is None else responses_config.auto_reasoning_summary() + + +def opt_in_auto_reasoning_summary(model: str | None, custom_llm_provider: object) -> ReasoningSummaryLevel | None: + if not (litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"): + return None + return "detailed" if model is None else auto_reasoning_summary(model, custom_llm_provider) + + +def reasoning_with_summary_default( + reasoning: Mapping[str, object], summary_alias: object, model: str, custom_llm_provider: object +) -> Mapping[str, object]: + if "summary" in reasoning: + return {**reasoning} + summary: Final = ( + summary_alias if summary_alias is not None else opt_in_auto_reasoning_summary(model, custom_llm_provider) + ) + return {**reasoning} if summary is None else {**reasoning, "summary": summary} + + +_SUMMARY_VALUES: Final = frozenset((*get_args(ReasoningSummaryLevel), None)) + + +def _openai_model_info(model: str) -> ModelInfo | None: + from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration + + model_info: Final = lookup_model_info_for_registration(model.removeprefix("responses/"), "openai") + return None if isinstance(model_info, ModelInfoLookupMiss) else model_info + + +def _declared_values(values: Sequence[str] | None) -> frozenset[str | None] | None: + return None if values is None else frozenset(values) + + +def _openai_reasoning_field_values(model: str) -> Mapping[str, frozenset[str | None] | None]: + model_info: Final = _openai_model_info(model) + if model_info is None: + return MappingProxyType({"summary": _SUMMARY_VALUES, "mode": None, "context": None}) + if not model_info.get("supports_reasoning"): + return MappingProxyType({}) + return MappingProxyType( + { + "summary": _SUMMARY_VALUES, + "mode": _declared_values(model_info.get("reasoning_modes")), + "context": _declared_values(model_info.get("reasoning_contexts")), + } + ) + + +def _openai_accepts_reasoning_field( + key: str, value: object, values_by_field: Mapping[str, frozenset[str | None] | None] +) -> bool: + if key == "effort": + return True + if key not in values_by_field: + return False + allowed: Final = values_by_field[key] + return allowed is None or (isinstance(value, str | None) and value in allowed) + + +def _openai_accepted_reasoning(reasoning: Mapping[str, object], model: str) -> Mapping[str, object]: + values_by_field: Final = _openai_reasoning_field_values(model) + return { + key: value for key, value in reasoning.items() if _openai_accepts_reasoning_field(key, value, values_by_field) + } + + +def _additional_drop_paths(additional_drop_params: object) -> tuple[str, ...]: + if additional_drop_params.__class__ is not list: + return () + return tuple(path for path in iter_items(additional_drop_params) if isinstance(path, str)) + + +def openai_bridge_reasoning( + *, + requested: Mapping[str, object], + effort_kept: bool, + summary_alias: object, + model: str, + drop_params: bool, + additional_drop_params: object, +) -> Mapping[str, object] | None: + drop_paths: Final = _additional_drop_paths(additional_drop_params) + if "reasoning_effort" in drop_paths: + return None + kept: Final = requested if effort_kept else {key: value for key, value in requested.items() if key != "effort"} + with_summary: Final = reasoning_with_summary_default(kept, summary_alias, model, "openai") + accepted: Final = ( + _openai_accepted_reasoning(with_summary, model) + if drop_params and OpenAIGPTConfig.is_openai_catalog_model(model) + else {**with_summary} + ) + remaining: Final = as_str_mapping( + _apply_nested_drop_params({"reasoning_effort": accepted}, list(drop_paths)).get("reasoning_effort") + ) + return remaining or None diff --git a/litellm/router.py b/litellm/router.py index e8d83e1203e..fc888495e99 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7,6 +7,29 @@ # # Thank you ! We ❤️ you! - Krrish & Ishaan +""" +litellm.Router: load balancing, retries, fallbacks and cooldowns across deployments. + +Public surface. The router is moving behind a narrow interface; code outside +the router should end up using a Router through two kinds of members only: + +- endpoint methods (acompletion, aembedding, aresponses, aanthropic_messages, + the files, batches, assistants and vector store wrappers, ...), each a thin + wrapper around one call; +- entry points (being added) that replace direct access to internals: registry() for + deployment lookups, settings() for configuration reads, health() for + cooldowns, blocks and usage, select() to pick a deployment without calling + it, apply() for runtime changes, observe() for routing events, + register_plugin() for strategies and sub-routers, and retire() to release + the router. + +Everything else on the class is internal. The members outside code uses today +are listed in tests/code_coverage_tests/router_surface_allowlist.json, and a CI +check (router_surface_check.py next to it) fails on any new outside use, so the +list only shrinks while callers move to the entry points. The same interface +is documented for the Rust core in litellm-rust/crates/router-core. +""" + import asyncio import contextlib import copy @@ -85,6 +108,15 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.routing import routing_decision_attributes from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.batch_utils import ( + batch_create_routing, + deployment_supported_endpoints, + get_router_metadata_variable_name, + is_batch_retrieve_call_type, + replace_model_in_jsonl, + should_replace_model_in_jsonl, + without_deployments_that_cannot_batch, +) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( coerce_token_limit, @@ -148,6 +180,15 @@ from litellm.llms.openai_like.model_info import ( get_openai_compatible_model_info, ) from litellm.responses.bridge_check import keep_registered_mode +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + AUTO_ROUTER_MODEL_PREFIX, + STRATEGY_ROUTER_PARAM_FIELDS, + GatedAutoRouterCapability, + capability_limit_violation, + claimed_capabilities, + classify_strategy_router_model, + count_capability_routers, +) from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.complexity_router.context_compaction import ( @@ -187,24 +228,6 @@ from litellm.router_utils.add_retry_fallback_headers import ( replace_complexity_router_headers, response_total_token_count, ) -from litellm.router_utils.auto_router_model_naming import ( - AUTO_ROUTER_MODEL_PREFIX, - STRATEGY_ROUTER_PARAM_FIELDS, - GatedAutoRouterCapability, - capability_limit_violation, - claimed_capabilities, - classify_strategy_router_model, - count_capability_routers, -) -from litellm.router_utils.batch_utils import ( - batch_create_routing, - deployment_supported_endpoints, - get_router_metadata_variable_name, - is_batch_retrieve_call_type, - replace_model_in_jsonl, - should_replace_model_in_jsonl, - without_deployments_that_cannot_batch, -) from litellm.router_utils.client_initalization_utils import InitalizeCachedClient, MaxParallelRequestsLimit from litellm.router_utils.clientside_credential_handler import ( get_dynamic_litellm_params, @@ -359,6 +382,7 @@ from litellm.router_utils.request_priority import ( InvalidPriority, parse_default_priority, priority_error, + priority_header_fields, request_drops_params, resolve_request_priority, unqueueable_priority_message, @@ -1369,7 +1393,7 @@ class Router: }, ] - router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": "openai-gpt-3.5-turbo"}]) + router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}]) ``` """ @@ -1595,6 +1619,10 @@ class Router: if default_fallbacks is not None or litellm.default_fallbacks is not None: _fallbacks = default_fallbacks or litellm.default_fallbacks + if isinstance(_fallbacks, str): + raise TypeError( + f"default_fallbacks must be a list of model names, got '{_fallbacks}'. Use ['{_fallbacks}']" + ) if self.fallbacks is not None: self.fallbacks.append({"*": _fallbacks}) else: @@ -2838,7 +2866,7 @@ class Router: self._initialize_core_endpoints() self._initialize_specialized_endpoints() - def validate_fallbacks(self, fallback_param: list | None): + def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None): """ Validate the fallbacks parameter. """ @@ -2851,6 +2879,13 @@ class Router: raise ValueError( f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." ) + ((model_group, targets),) = fallback_dict.items() + if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=[fallback_dict]): + continue + raise TypeError( + f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. " + f"Use {{'{model_group}': ['{targets}']}}" + ) def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None: def _move_before_deployment_affinity( @@ -5411,14 +5446,22 @@ class Router: request_priority: Final = resolve_request_priority( requested=kwargs.get("priority"), default_priority=self.default_priority, - drops_params=lambda: request_drops_params( - kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs) - ), + drops_params=lambda: self._drops_params(model=model, kwargs=kwargs), ) if isinstance(request_priority, InvalidPriority): raise priority_error(model=model, message=request_priority.message) return request_priority, {key: value for key, value in kwargs.items() if key != "priority"} + def priority_header_update( + self, model: str, kwargs: Mapping[str, object], header: str | None + ) -> Mapping[str, int | str]: + return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs)) + + def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool: + return request_drops_params( + kwargs, self.default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs) + ) + async def _schedule_factory( self, model: str, @@ -6456,6 +6499,9 @@ class Router: hop_kwargs: Final = mid_stream_fallback_hop_kwargs( model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs ) + shared_logging: Final = kwargs.get("litellm_logging_obj") + if isinstance(shared_logging, LiteLLMLogging): + shared_logging.carry_failed_attempt_usage() response: Final = await self._ageneric_api_call_with_fallbacks_helper( model=model, original_generic_function=original_generic_function, **kwargs ) @@ -10364,28 +10410,38 @@ class Router: model_info[field] = backend_value @staticmethod - def _cost_map_backend_model(deployment: Deployment) -> str: - model_info_base_model: Final = deployment.model_info.base_model + def _cost_map_backend_model(deployment: Deployment | DeploymentTypedDict) -> str: + model_info_base_model: Final = ( + deployment.model_info.base_model + if isinstance(deployment, Deployment) + else _configured_model_info(deployment).get("base_model") + ) if isinstance(model_info_base_model, str) and model_info_base_model: return model_info_base_model - params_base_model: Final = deployment.litellm_params.get("base_model") + params: Final = _LITELLM_PARAMS_ADAPTER.validate_python( + deployment.litellm_params.model_dump(exclude_none=True) + if isinstance(deployment, Deployment) + else deployment["litellm_params"] + ) + params_base_model: Final = params.get("base_model") if isinstance(params_base_model, str) and params_base_model: return params_base_model - params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) + model: Final = ( + deployment.litellm_params.model + if isinstance(deployment, Deployment) + else deployment["litellm_params"].get("model", "") + ) invoked_model: Final = params.get("model_id") if ( isinstance(invoked_model, str) and invoked_model - and ( - deployment.litellm_params.custom_llm_provider == "bedrock" - or deployment.litellm_params.model.startswith("bedrock/") - ) + and (params.get("custom_llm_provider") == "bedrock" or model.startswith("bedrock/")) ): profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}" profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock") if not isinstance(profile_info, ModelInfoLookupMiss): return profile_model - return deployment.litellm_params.model + return model @staticmethod def _inherit_builtin_service_tier_pricing( @@ -11837,6 +11893,11 @@ class Router: indices: Final = self.model_name_to_deployment_indices.get(model_name) return None if indices is None else tuple(self.model_list[idx] for idx in indices) + def team_public_model_names_of(self, team_id: str) -> frozenset[str]: + """The public names `team_id` serves its own deployments under.""" + with self._deployment_index_lock: + return frozenset(name for owner, name in self.team_model_to_deployment_indices if owner == team_id) + def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None: """The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed.""" with self._deployment_index_lock: @@ -13945,7 +14006,9 @@ class Router: ] return named or self._get_wildcard_deployments(model_name=model_group) - def resolved_litellm_models(self, model_name: str, team_id: str | None = None) -> tuple[str, ...]: + def resolved_litellm_models( + self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False + ) -> tuple[str, ...]: """The provider model strings `model_name` can actually be served by on this proxy. `get_model_list` composes every channel the request path itself uses (exact name, @@ -13958,6 +14021,12 @@ class Router: deployment behind it - so the fallback for an empty result is the caller's policy, never this function's. """ + if use_base_model: + return tuple( + backend_model + for deployment in self.get_model_list(model_name=model_name, team_id=team_id) or () + if (backend_model := self._cost_map_backend_model(deployment)) + ) return tuple( litellm_model for deployment in self.get_model_list(model_name=model_name, team_id=team_id) or () @@ -16225,7 +16294,7 @@ class Router: to the deployment that actually served the request. Every attempt therefore writes or clears, never just writes. """ - from litellm.router_utils.baseline_request import capture_baseline_parameters + from litellm.router_strategy.auto_router.baseline_request import capture_baseline_parameters from litellm.types.router import BaselineRouteStamp phase_attributes(routing_decision_attributes(routing_decision)) diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_strategy/auto_router/auto_router_model_naming.py similarity index 100% rename from litellm/router_utils/auto_router_model_naming.py rename to litellm/router_strategy/auto_router/auto_router_model_naming.py diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_strategy/auto_router/auto_router_tuning_baseline.py similarity index 100% rename from litellm/router_utils/auto_router_tuning_baseline.py rename to litellm/router_strategy/auto_router/auto_router_tuning_baseline.py diff --git a/litellm/router_utils/baseline_request.py b/litellm/router_strategy/auto_router/baseline_request.py similarity index 100% rename from litellm/router_utils/baseline_request.py rename to litellm/router_strategy/auto_router/baseline_request.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index eab232eae18..1d502dd8d95 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -36,12 +36,12 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs -from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.router_utils.cooldown_callbacks import ( get_prometheus_logger_from_callbacks, ) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 6f3e0936641..8e0c3d6a7e7 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 5894da9cb4f..17d03fd4ff8 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -10,7 +10,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 4e994d3fdc6..c352934db6c 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -13,9 +13,9 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, safe_divide_seconds from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import LiteLLMPydanticObjectBase diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 2d373e0c266..9636d7f009d 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -9,7 +9,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 527ea830d55..0719b30d36d 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -15,8 +15,8 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 128bca6cca0..a1b1e327802 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -13,7 +13,7 @@ from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm._internal_context import service_target @@ -417,17 +417,35 @@ def _should_run_cooldown_logic( return True -_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, list[str | Mapping[str, object]]] | str, ...]) +_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[object, ...]) +_FALLBACK_ENTRY_ADAPTER: Final = TypeAdapter(Mapping[str, list[str | Mapping[str, object]] | str]) def _fallback_target_name(target: str | Mapping[str, object]) -> str: return target if isinstance(target, str) else str(target.get("model", "")) +def _fallback_target_names( + targets: Sequence[str | Mapping[str, object]] | str, +) -> list[str]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + return [targets] if isinstance(targets, str) else [_fallback_target_name(target) for target in targets] + + def _fallback_chain_names( - entry: Mapping[str, Sequence[str | Mapping[str, object]]], -) -> Mapping[str, list[str]]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape - return {key: [_fallback_target_name(target) for target in targets] for key, targets in entry.items()} + entry: object, +) -> Mapping[str, list[str]] | None: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + try: + validated: Final = _FALLBACK_ENTRY_ADAPTER.validate_python(entry) + except ValidationError: + return None + return {key: _fallback_target_names(targets) for key, targets in validated.items()} + + +def _fallback_entries(configured: object) -> tuple[object, ...]: + try: + return _FALLBACK_ENTRIES_ADAPTER.validate_python(configured) + except ValidationError: + return () def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool: @@ -437,9 +455,7 @@ def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_ if requested_model_group is None or not configured: return False chains: Final = tuple( - _fallback_chain_names(entry) - for entry in _FALLBACK_ENTRIES_ADAPTER.validate_python(configured) - if isinstance(entry, Mapping) + chain for chain in map(_fallback_chain_names, _fallback_entries(configured)) if chain is not None ) fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group) return bool(fallback_group) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 0a77ad5b9ea..6d1de1f9f6c 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -13,6 +13,7 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_router_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider @@ -22,7 +23,6 @@ from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, get_fallback_error_info, ) -from litellm.router_utils.batch_utils import get_router_metadata_variable_name from litellm.router_utils.cooldown_handlers import ( _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper, used across router_utils cast_exception_status_to_int, diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index ea3fdd601dc..9eebcf39515 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -5,15 +5,19 @@ When enabled via router_settings.optional_pre_call_checks: ["enforce_model_rate_ - tpm/rpm: combined TPM + optional RPM (legacy) - itpm/otpm: separate input/output tokens per minute +- tpd/rpd: tokens and requests per UTC calendar day When a deployment sets both itpm/otpm and tpm/rpm, both are enforced. A warning is logged the first time such a deployment is seen. """ import contextlib -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Callable +from datetime import datetime, timedelta +from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from pydantic import BaseModel, ConfigDict import litellm from litellm._internal_context import with_service_target @@ -33,6 +37,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( io_token_refund_failure, ) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.router import RouterErrors from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_utc_datetime @@ -47,6 +52,62 @@ else: class RoutingArgs: ttl: int = 60 # 1min (RPM/TPM expire key) + daily_ttl: int = 86400 + + +DAY_FORMAT: Final = "%Y-%m-%d" + + +def seconds_until_utc_midnight(now: datetime) -> int: + next_midnight: Final = (now + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0) + return max(1, int((next_midnight - now).total_seconds())) + + +def _first_set(*values: int | None) -> int | None: + return next((value for value in values if value is not None), None) + + +class _DailyLimitFields(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True) + tpd: int | None = None + rpd: int | None = None + + +class _DailyLimitParams(_DailyLimitFields): + model: str | None = None + + +class _DailyLimitModelInfo(_DailyLimitFields): + id: str | None = None + + +class _DailyLimitDeployment(_DailyLimitFields): + model_name: str = "" + litellm_params: _DailyLimitParams = _DailyLimitParams() + model_info: _DailyLimitModelInfo = _DailyLimitModelInfo() + + def limits(self) -> tuple[int | None, int | None]: + return ( + _first_set(self.tpd, self.litellm_params.tpd, self.model_info.tpd), + _first_set(self.rpd, self.litellm_params.rpd, self.model_info.rpd), + ) + + def cache_keys(self, current_day: str) -> tuple[str, str]: + prefix: Final = f"{self.model_info.id}:{self.litellm_params.model}" + return f"{prefix}:tpd:{current_day}", f"{prefix}:rpd:{current_day}" + + def limit_error(self, kind: Literal["tpd", "rpd"], limit: int, usage: int, now: datetime) -> litellm.RateLimitError: + return litellm.RateLimitError( + message=f"Model rate limit exceeded. {kind.upper()} limit={limit}, current usage={usage}", + llm_provider="", + model=self.litellm_params.model, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} {kind} limit={limit}. current usage={usage}. id={self.model_info.id}, model_group={self.model_name}", + headers={"retry-after": str(seconds_until_utc_midnight(now))}, + request=httpx.Request(method="model_rate_limit_check", url="https://github.com/BerriAI/litellm"), + ), + ) class ModelRateLimitingCheck(CustomLogger): @@ -60,12 +121,60 @@ class ModelRateLimitingCheck(CustomLogger): this check actively enforces those limits across ALL routing strategies. """ - def __init__(self, dual_cache: DualCache): + def __init__(self, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime): self.dual_cache = dual_cache + self.clock = clock # model_ids already warned about itpm/otpm + tpm/rpm on the same deployment, # so the warning is logged once per deployment rather than per request. self._io_token_conflict_warned_ids: set[str] = set() + def _increment(self, key: str, value: int, ttl: int) -> int: + return self.dual_cache.increment_cache(key=key, value=value, ttl=ttl) + + async def _async_increment( + self, key: str, value: int, ttl: int, parent_otel_span: Span | None = None + ) -> int | None: + count: Final = await self.dual_cache.async_increment_cache( + key=key, value=value, ttl=ttl, parent_otel_span=parent_otel_span + ) + return None if count is None else int(count) + + @staticmethod + def _request_count_operations( + rpm_key: str | None, rpd_key: str | None, increment_value: int + ) -> tuple[RedisPipelineIncrementOperation, ...]: + return tuple( + RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=ttl) + for key, ttl in ((rpm_key, RoutingArgs.ttl), (rpd_key, RoutingArgs.daily_ttl)) + if key is not None + ) + + async def _async_refund_request_counts( + self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None + ) -> None: + operations: Final = self._request_count_operations(rpm_key, rpd_key, -1) + if operations: + await self.dual_cache.async_increment_cache_pipeline( + increment_list=list(operations), parent_otel_span=parent_otel_span + ) + + async def _async_increment_request_counts( + self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None + ) -> tuple[int | None, int | None]: + operations: Final = self._request_count_operations(rpm_key, rpd_key, 1) + if not operations: + return None, None + counts: Final = await self.dual_cache.async_increment_cache_pipeline( + increment_list=list(operations), parent_otel_span=parent_otel_span + ) + if counts is None or len(counts) != len(operations): + return None, None + count_by_key: Final = {op["key"]: int(count) for op, count in zip(operations, counts)} + return ( + None if rpm_key is None else count_by_key.get(rpm_key), + None if rpd_key is None else count_by_key.get(rpd_key), + ) + def _warn_io_token_and_tpm_rpm_coexist_once(self, deployment: dict) -> None: tpm_limit, rpm_limit = self._get_deployment_limits(deployment) if tpm_limit is None and rpm_limit is None: @@ -129,6 +238,21 @@ class ModelRateLimitingCheck(CustomLogger): return tpm, rpm + def _check_daily_limits(self, deployment: _DailyLimitDeployment, now: datetime) -> None: + tpd_limit, rpd_limit = deployment.limits() + if tpd_limit is None and rpd_limit is None: + return + tpd_key, rpd_key = deployment.cache_keys(now.strftime(DAY_FORMAT)) + if tpd_limit is not None: + current_tpd: Final = self._get_current_tpm(tpd_key, tpd_limit) + if current_tpd is not None and current_tpd >= tpd_limit: + raise deployment.limit_error("tpd", tpd_limit, current_tpd, now) + if rpd_limit is not None: + current_rpd: Final = self._increment(rpd_key, 1, RoutingArgs.daily_ttl) + if current_rpd > rpd_limit: + self._increment(rpd_key, -1, RoutingArgs.daily_ttl) + raise deployment.limit_error("rpd", rpd_limit, current_rpd, now) + def _get_cache_keys(self, deployment: dict, current_minute: str) -> tuple[str, str]: """Get the cache keys for TPM and RPM tracking.""" model_id: Final = deployment.get("model_info", {}).get("id") @@ -178,13 +302,10 @@ class ModelRateLimitingCheck(CustomLogger): ) io_reservation_made = True + dt: Final = self.clock() + daily: Final = _DailyLimitDeployment.model_validate(deployment) tpm_limit, rpm_limit = self._get_deployment_limits(deployment) - # If no limits are set, allow the request - if tpm_limit is None and rpm_limit is None: - return deployment - - dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) @@ -213,8 +334,9 @@ class ModelRateLimitingCheck(CustomLogger): # Check RPM limit (atomic increment-first to avoid race conditions) if rpm_limit is not None: - current_rpm: Final = self.dual_cache.increment_cache(key=rpm_key, value=1, ttl=RoutingArgs.ttl) + current_rpm: Final = self._increment(rpm_key, 1, RoutingArgs.ttl) if current_rpm is not None and current_rpm > rpm_limit: + self._increment(rpm_key, -1, RoutingArgs.ttl) raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", llm_provider="", @@ -230,6 +352,12 @@ class ModelRateLimitingCheck(CustomLogger): ), ) + try: + self._check_daily_limits(daily, dt) + except litellm.RateLimitError: + if rpm_limit is not None: + self._increment(rpm_key, -1, RoutingArgs.ttl) + raise return deployment except litellm.RateLimitError: @@ -259,15 +387,14 @@ class ModelRateLimitingCheck(CustomLogger): ) io_reservation_made = True + dt: Final = self.clock() + daily: Final = _DailyLimitDeployment.model_validate(deployment) + tpd_limit, rpd_limit = daily.limits() tpm_limit, rpm_limit = self._get_deployment_limits(deployment) - # If no limits are set, allow the request - if tpm_limit is None and rpm_limit is None: - return deployment - - dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + tpd_key, rpd_key = daily.cache_keys(dt.strftime(DAY_FORMAT)) model_id: Final = deployment.get("model_info", {}).get("id") model_name: Final = deployment.get("litellm_params", {}).get("model") @@ -293,30 +420,37 @@ class ModelRateLimitingCheck(CustomLogger): num_retries=0, # Don't retry - return 429 immediately ) - # Check RPM limit (atomic increment-first to avoid race conditions) - if rpm_limit is not None: - current_rpm: Final = await self.dual_cache.async_increment_cache( - key=rpm_key, - value=1, - ttl=RoutingArgs.ttl, - parent_otel_span=parent_otel_span, - ) - if current_rpm is not None and current_rpm > rpm_limit: - raise litellm.RateLimitError( - message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", - llm_provider="", - model=model_name, - response=httpx.Response( - status_code=429, - content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", - headers={"retry-after": str(60)}, - request=httpx.Request( - method="model_rate_limit_check", - url="https://github.com/BerriAI/litellm", - ), + if tpd_limit is not None: + current_tpd: Final = await self._async_get_current_tpm(tpd_key, tpd_limit, parent_otel_span) + if current_tpd is not None and current_tpd >= tpd_limit: + raise daily.limit_error("tpd", tpd_limit, current_tpd, dt) + + charged_rpm_key: Final = rpm_key if rpm_limit is not None else None + charged_rpd_key: Final = rpd_key if rpd_limit is not None else None + current_rpm, current_rpd = await self._async_increment_request_counts( + charged_rpm_key, charged_rpd_key, parent_otel_span + ) + if rpm_limit is not None and current_rpm is not None and current_rpm > rpm_limit: + await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) + raise litellm.RateLimitError( + message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", + llm_provider="", + model=model_name, + response=httpx.Response( + status_code=429, + content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}", + headers={"retry-after": str(60)}, + request=httpx.Request( + method="model_rate_limit_check", + url="https://github.com/BerriAI/litellm", ), - num_retries=0, # Don't retry - return 429 immediately - ) + ), + num_retries=0, # Don't retry - return 429 immediately + ) + + if rpd_limit is not None and current_rpd is not None and current_rpd > rpd_limit: + await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) + raise daily.limit_error("rpd", rpd_limit, current_rpd, dt) return deployment @@ -373,17 +507,16 @@ class ModelRateLimitingCheck(CustomLogger): if not model or not total_tokens: return - dt: Final = get_utc_datetime() + dt: Final = self.clock() current_minute: Final = dt.strftime("%H-%M") tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" + tpd_key: Final = f"{model_id}:{model}:tpd:{dt.strftime(DAY_FORMAT)}" verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm_key, total_tokens) - await self.dual_cache.async_increment_cache( - key=tpm_key, - value=total_tokens, - ttl=RoutingArgs.ttl, - ) + tokens: Final = int(total_tokens) + await self._async_increment(tpm_key, tokens, RoutingArgs.ttl) + await self._async_increment(tpd_key, tokens, RoutingArgs.daily_ttl) except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.async_log_success_event: %s", e) @@ -434,15 +567,14 @@ class ModelRateLimitingCheck(CustomLogger): if not model or not total_tokens: return - dt: Final = get_utc_datetime() + dt: Final = self.clock() current_minute: Final = dt.strftime("%H-%M") tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" + tpd_key: Final = f"{model_id}:{model}:tpd:{dt.strftime(DAY_FORMAT)}" - self.dual_cache.increment_cache( - key=tpm_key, - value=total_tokens, - ttl=RoutingArgs.ttl, - ) + tokens: Final = int(total_tokens) + self._increment(tpm_key, tokens, RoutingArgs.ttl) + self._increment(tpd_key, tokens, RoutingArgs.daily_ttl) except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.log_success_event: %s", e) diff --git a/litellm/router_utils/request_priority.py b/litellm/router_utils/request_priority.py index 14608cc1bcf..61627d4b0cf 100644 --- a/litellm/router_utils/request_priority.py +++ b/litellm/router_utils/request_priority.py @@ -47,6 +47,15 @@ def resolve_request_priority( return priority +def priority_header_fields(header: str | None, drops_params: Callable[[], bool]) -> Mapping[str, int | str]: + if not header: + return {} + priority: Final = _as_priority(header) + if isinstance(priority, InvalidPriority): + return {} if drops_params() else {"priority": header} + return {"priority": priority} + + def parse_default_priority(value: object) -> int | None | InvalidPriority: if value is None: return None diff --git a/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index c197b21bad0..c2a030ca4d9 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -46,6 +46,12 @@ - A failure after execution starts is terminal and never triggers fallback - Creating an unstarted native coroutine acquires no clients or credentials and captures no execution context +# Native extension + +- `_native` ships in every wheel and the sdist builds it, so it is always importable +- Import it with `from litellm.rust_bridge import _native` and use the attributes `_native.pyi` declares +- Never add a `None` branch, `getattr` default or stand-in class for a missing extension + # Verification - Check cross-language behavior against a freshly installed extension, covering unstarted-call release, setup ordering, exception identity, sync and async streams, cancellation and GC diff --git a/litellm/rust_bridge/__init__.py b/litellm/rust_bridge/__init__.py index 8f6f4390b8a..52cc3c8cff1 100644 --- a/litellm/rust_bridge/__init__.py +++ b/litellm/rust_bridge/__init__.py @@ -1,10 +1,5 @@ """LiteLLM Rust bridge package.""" from litellm.rust_bridge.configuration import rust -from litellm.rust_bridge.loader import ( - get_native_bridge, - native_bridge_available, - reset_native_bridge_cache, -) -__all__ = ["get_native_bridge", "native_bridge_available", "reset_native_bridge_cache", "rust"] +__all__ = ["rust"] diff --git a/litellm/rust_bridge/bindings.py b/litellm/rust_bridge/bindings.py index d16f150a2aa..91570c6a8bc 100644 --- a/litellm/rust_bridge/bindings.py +++ b/litellm/rust_bridge/bindings.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Callable from typing import Final, Generic, TypeVar -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge import _native BindingT = TypeVar("BindingT") @@ -18,32 +18,24 @@ _UNSET: Final = _Unset() class NativeBinding(Generic[BindingT]): """Resolve one native attribute with an explicit, resettable test override.""" - def __init__(self, attribute: str, *, validate: Callable[[object], BindingT | None]) -> None: + def __init__( + self, attribute: str, *, validate: Callable[[object], BindingT | None], source: object = _native + ) -> None: self._attribute: Final = attribute self._validate: Final = validate + self._source: Final = source self._override: BindingT | None | _Unset = _UNSET def load(self) -> BindingT | None: if not isinstance(self._override, _Unset): return self._override - native: Final = get_native_bridge() - if native is None: - return None - return self._validate(getattr(native, self._attribute, None)) + return self._validate(getattr(self._source, self._attribute, None)) + + def with_source(self, source: object) -> NativeBinding[BindingT]: + return NativeBinding(self._attribute, validate=self._validate, source=source) def override(self, value: BindingT | None) -> None: self._override = value def reset(self) -> None: self._override = _UNSET - - -def native_exception_types() -> tuple[type[BaseException], type[BaseException]] | None: - native: Final = get_native_bridge() - if native is None: - return None - declined: Final = getattr(native, "RustBridgeDeclined", None) - upstream: Final = getattr(native, "RustUpstreamError", None) - if not isinstance(declined, type) or not isinstance(upstream, type): - return None - return declined, upstream diff --git a/litellm/rust_bridge/clickhouse.py b/litellm/rust_bridge/clickhouse.py index cf2c8fdd3e2..fe2809ef856 100644 --- a/litellm/rust_bridge/clickhouse.py +++ b/litellm/rust_bridge/clickhouse.py @@ -1,12 +1,10 @@ import os -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Final, Protocol, runtime_checkable - -from pydantic import ConfigDict, TypeAdapter +from typing import Final from litellm.constants import DEFAULT_AGENT_TRACING_RETENTION_DAYS, DEFAULT_CLICKHOUSE_DATABASE -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge import _native @dataclass(frozen=True, slots=True, repr=False) @@ -16,27 +14,6 @@ class SpendStorageConfig: retention_days: int = DEFAULT_AGENT_TRACING_RETENTION_DAYS -class NativeSpendConfig(Protocol): - def __init__(self, database: str, url: str, retention_days: int) -> None: ... - - -class NativeSpendStorage(Protocol): - def __init__(self, config: NativeSpendConfig) -> None: ... - def ensure_schema(self) -> Awaitable[None]: ... - def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Awaitable[None]: ... - - -@runtime_checkable -class NativeSpendBridge(Protocol): - NativeClickHouseSpendConfig: type[NativeSpendConfig] - NativeClickHouseSpendStorage: type[NativeSpendStorage] - - -_NATIVE: Final[TypeAdapter[NativeSpendBridge]] = TypeAdapter( - NativeSpendBridge, config=ConfigDict(arbitrary_types_allowed=True) -) - - def spend_storage_config(environ: Mapping[str, str] = os.environ) -> SpendStorageConfig: url: Final = environ.get("CLICKHOUSE_URL") if not url: @@ -53,12 +30,8 @@ def spend_storage_config(environ: Mapping[str, str] = os.environ) -> SpendStorag class ClickHouseSpendStorage: def __init__(self, config: SpendStorageConfig) -> None: - bridge: Final = get_native_bridge() - if bridge is None: - raise RuntimeError("ClickHouse spend logging requires the Rust extension") - native: Final = _NATIVE.validate_python(bridge) - self._native: Final = native.NativeClickHouseSpendStorage( - native.NativeClickHouseSpendConfig(config.database, config.url, config.retention_days) + self._native: Final = _native.NativeClickHouseSpendStorage( + _native.NativeClickHouseSpendConfig(config.database, config.url, config.retention_days) ) async def ensure_schema(self) -> None: diff --git a/litellm/rust_bridge/fork_guard.py b/litellm/rust_bridge/fork_guard.py index c94665fb8db..82eb5c96295 100644 --- a/litellm/rust_bridge/fork_guard.py +++ b/litellm/rust_bridge/fork_guard.py @@ -11,32 +11,22 @@ site, so the runtime can never start there. Workers forked from it are unaffecte from __future__ import annotations +from collections.abc import Callable from typing import Final -from litellm.rust_bridge.loader import get_native_bridge +from litellm.rust_bridge import _native class NativeStateStartedBeforeFork(RuntimeError): pass -class _NeverRaised(RuntimeError): - """Stands in for a native exception when the extension is unavailable or predates it.""" +ForkedAfterNativeRuntimeStarted: Final = _native.ForkedAfterNativeRuntimeStarted +ProcessReservedForForking: Final = _native.ProcessReservedForForking -_native: Final = get_native_bridge() -ForkedAfterNativeRuntimeStarted: Final[type[RuntimeError]] = getattr( - _native, "ForkedAfterNativeRuntimeStarted", _NeverRaised -) -ProcessReservedForForking: Final[type[RuntimeError]] = getattr(_native, "ProcessReservedForForking", _NeverRaised) - - -def reserve_process_for_forking(where: str) -> None: +def reserve_process_for_forking(where: str, reserve: Callable[[], None] = _native.reserve_process_for_forking) -> None: """Forbid native routes in this process. Raises if one already ran here.""" - native: Final = get_native_bridge() - reserve: Final = getattr(native, "reserve_process_for_forking", None) - if not callable(reserve): - return try: reserve() except RuntimeError as error: diff --git a/litellm/rust_bridge/loader.py b/litellm/rust_bridge/loader.py deleted file mode 100644 index 022c38f5a85..00000000000 --- a/litellm/rust_bridge/loader.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Loader for the packaged LiteLLM Rust extension.""" - -from __future__ import annotations - -from types import ModuleType -from typing import Final - -_BRIDGE_SENTINEL: Final = object() -_cached_bridge: ModuleType | None | object = _BRIDGE_SENTINEL - - -def get_native_bridge() -> ModuleType | None: - """Return the packaged Rust extension, or ``None`` when unavailable.""" - global _cached_bridge - if _cached_bridge is not _BRIDGE_SENTINEL: - return _cached_bridge if isinstance(_cached_bridge, ModuleType) else None - - try: - from litellm.rust_bridge import _native - except ImportError: - _cached_bridge = None - return None - _cached_bridge = _native - return _native - - -def reset_native_bridge_cache() -> None: - """Forget the cached extension so the next lookup reimports it from disk.""" - global _cached_bridge - _cached_bridge = _BRIDGE_SENTINEL - - -def native_bridge_available() -> bool: - """Whether the packaged Rust extension is importable.""" - return get_native_bridge() is not None diff --git a/litellm/rust_bridge/responses/websocket.py b/litellm/rust_bridge/responses/websocket.py index f10631958bb..83a1fa53376 100644 --- a/litellm/rust_bridge/responses/websocket.py +++ b/litellm/rust_bridge/responses/websocket.py @@ -2,32 +2,33 @@ from __future__ import annotations +from collections.abc import Awaitable from dataclasses import dataclass from typing import Final, Protocol import httpx from websockets.exceptions import ConnectionClosedOK +from litellm.rust_bridge import _native from litellm.rust_bridge.host.timeouts import timeout_to_seconds -from litellm.rust_bridge.loader import get_native_bridge class RustResponsesWebSocket(Protocol): - async def send_text(self, text: str) -> None: ... + def send_text(self, text: str) -> Awaitable[None]: ... - async def recv_text(self) -> str | None: ... + def recv_text(self) -> Awaitable[str | None]: ... - async def close(self) -> None: ... + def close(self) -> Awaitable[None]: ... class RustResponsesWebSocketConnection(Protocol): @classmethod - async def connect( + def connect( cls, url: str, headers: dict[str, str], timeout_seconds: float | None, - ) -> RustResponsesWebSocket: ... + ) -> Awaitable[RustResponsesWebSocket]: ... class _Unset: @@ -53,16 +54,10 @@ def set_rust_responses_websocket( _STATE.connection = connection -def load_rust_responses_websocket() -> RustResponsesWebSocketConnection | None: +def load_rust_responses_websocket() -> RustResponsesWebSocketConnection: if _STATE.connection is not None: return _STATE.connection - native_bridge: Final = get_native_bridge() - if native_bridge is None: - return None - connection_type: Final[RustResponsesWebSocketConnection | None] = getattr( - native_bridge, "ResponsesWebSocketConnection", None - ) - return connection_type + return _native.ResponsesWebSocketConnection class _ConnectionAdapter: @@ -89,8 +84,6 @@ async def connect( timeout: float | httpx.Timeout | None, ) -> _ConnectionAdapter | None: connection_type: Final = load_rust_responses_websocket() - if connection_type is None: - return None try: connection: Final = await connection_type.connect( url=url, diff --git a/litellm/rust_bridge/runtime.py b/litellm/rust_bridge/runtime.py index cf632c09074..e22ebe73b08 100644 --- a/litellm/rust_bridge/runtime.py +++ b/litellm/rust_bridge/runtime.py @@ -7,7 +7,8 @@ from typing import Final, Generic, NoReturn, TypeAlias, TypeVar from typing_extensions import assert_never from litellm.exceptions import APIError -from litellm.rust_bridge.bindings import NativeBinding, native_exception_types +from litellm.rust_bridge import _native +from litellm.rust_bridge.bindings import NativeBinding from litellm.rust_bridge.catalog import RouteContext, Rules, decision from litellm.rust_bridge.configuration import Decision from litellm.rust_bridge.response_metadata import mark_rust_response @@ -153,15 +154,11 @@ def attempt( ) -> RustAttempt[ResultT]: if native_call is None: return RustUnavailable() - exceptions: Final = native_exception_types() - if exceptions is None: - return RustHandled(adapt(native_call())) - declined, upstream = exceptions try: value: Final = native_call() - except declined as error: + except _native.RustBridgeDeclined as error: return RustDeclined(reason=_decline_reason(error)) - except upstream as error: + except _native.RustUpstreamError as error: _raise_upstream(error, context) return RustHandled(adapt(value)) @@ -174,15 +171,11 @@ async def aattempt( ) -> RustAttempt[ResultT]: if native_call is None: return RustUnavailable() - exceptions: Final = native_exception_types() - if exceptions is None: - return RustHandled(adapt(await native_call())) - declined, upstream = exceptions try: value: Final = await native_call() - except declined as error: + except _native.RustBridgeDeclined as error: return RustDeclined(reason=_decline_reason(error)) - except upstream as error: + except _native.RustUpstreamError as error: _raise_upstream(error, context) return RustHandled(adapt(value)) diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 98d12e497da..de8566e6fdc 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -21,6 +21,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final import httpx +from pydantic import ConfigDict, TypeAdapter, with_config +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -41,6 +43,15 @@ if TYPE_CHECKING: from botocore.awsrequest import HTTPHeaders +@with_config(ConfigDict(extra="allow", strict=True)) +class _SecretStringResponse(TypedDict): + SecretString: ReadOnly[str] + + +_SECRET_STRING_RESPONSE: Final = TypeAdapter(_SecretStringResponse) +_AWS_SECRET_MANAGER_RESPONSE: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) + + class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): def __init__( self, @@ -154,7 +165,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json()["SecretString"] + return _SECRET_STRING_RESPONSE.validate_python(response.json())["SecretString"] except httpx.TimeoutException: raise ValueError("Timeout error occurred") except Exception as e: @@ -204,7 +215,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = sync_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) - return response.json()["SecretString"] + return _SECRET_STRING_RESPONSE.validate_python(response.json())["SecretString"] except httpx.TimeoutException: raise ValueError("Timeout error occurred") except httpx.HTTPStatusError as e: @@ -424,7 +435,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -481,7 +492,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -533,7 +544,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: @@ -613,7 +624,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): try: response: Final = await async_client.post(url=endpoint_url, headers=headers, data=body.decode("utf-8")) response.raise_for_status() - return response.json() + return _AWS_SECRET_MANAGER_RESPONSE.validate_python(response.json()) except httpx.HTTPStatusError as err: raise ValueError(f"HTTP error occurred: {err.response.text}") except httpx.TimeoutException: diff --git a/litellm/tracing/otlp_http.py b/litellm/tracing/otlp_http.py index 493cc7cf984..81dbab6c023 100644 --- a/litellm/tracing/otlp_http.py +++ b/litellm/tracing/otlp_http.py @@ -1,20 +1,9 @@ import json -from typing import Final, Protocol, runtime_checkable +from typing import Final -from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict -from litellm.rust_bridge.loader import get_native_bridge - - -@runtime_checkable -class NativeOtlpError(Protocol): - def trace_encode_error(self, message: str) -> bytes: ... - - -_NATIVE: Final[TypeAdapter[NativeOtlpError]] = TypeAdapter( - NativeOtlpError, config=ConfigDict(arbitrary_types_allowed=True) -) +from litellm.rust_bridge import _native class OTLPError(TypedDict): @@ -22,8 +11,7 @@ class OTLPError(TypedDict): def encode_error(message: str) -> bytes: - native: Final = get_native_bridge() - return b"" if native is None else _NATIVE.validate_python(native).trace_encode_error(message) + return _native.trace_encode_error(message) def encode_otlp_response(content_type: str | None, error: str | None = None) -> tuple[bytes, str]: diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 97dd67d3ff8..826aa5d9c98 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -3,7 +3,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias from urllib.parse import urlsplit -from pydantic import ConfigDict, Field, PrivateAttr, StrictInt, field_validator +from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, StrictInt, field_validator from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMBaseModel, LiteLLMPydanticObjectBase @@ -343,6 +343,18 @@ class ListAgentsResponse(LiteLLMBaseModel): agents: list[AgentResponse] +class AgentCreateRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + id: str + base_agent: str | None = None + agent_config: Mapping[str, object] | None = None + system_instruction: str | None = None + base_environment: str | Mapping[str, object] | None = None + description: str | None = None + tools: Sequence[Mapping[str, object]] | None = None + + class AgentCreateResponse(LiteLLMPydanticObjectBase): """ Response from a provider-side agent creation or get call (e.g. Gemini v1beta/agents). diff --git a/litellm/types/integrations/rag/bedrock_knowledgebase.py b/litellm/types/integrations/rag/bedrock_knowledgebase.py index 7156d8101e1..7288eca03b4 100644 --- a/litellm/types/integrations/rag/bedrock_knowledgebase.py +++ b/litellm/types/integrations/rag/bedrock_knowledgebase.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from typing_extensions import ReadOnly, TypedDict @@ -40,7 +40,7 @@ class BedrockKBRetrievalResult(TypedDict, total=False): content: BedrockKBContent | None location: BedrockKBLocation | None score: float | None - metadata: dict[str, Any] | None + metadata: dict[str, object] | None class BedrockKBResponse(TypedDict, total=False): @@ -87,21 +87,21 @@ class BedrockKBMetadataConfiguration(TypedDict, total=False): class BedrockKBModelConfiguration(TypedDict, total=False): """Model configuration for reranking.""" - additionalModelRequestFields: dict[str, Any] | None + additionalModelRequestFields: dict[str, object] | None modelArn: str | None class BedrockKBRerankingConfiguration(TypedDict, total=False): """Configuration for reranking in vector search.""" - bedrockRerankingConfiguration: dict[str, Any] | None # This could be further typed if needed + bedrockRerankingConfiguration: dict[str, object] | None # This could be further typed if needed type: str | None class BedrockKBVectorSearchConfiguration(TypedDict, total=False): """Configuration for vector search.""" - filter: dict[str, Any] | None + filter: dict[str, object] | None implicitFilterConfiguration: BedrockKBImplicitFilterConfiguration | None numberOfResults: int | None overrideSearchType: str | None diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 717b74e5741..9196b7a226f 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -189,6 +189,8 @@ class DeploymentOptions: tpm: int | None = None itpm: int | None = None otpm: int | None = None + rpd: int | None = None + tpd: int | None = None kubernetes_pod_discovery: bool | None = None default_api_key_rpm_limit: int | None = None default_api_key_tpm_limit: int | None = None diff --git a/litellm/types/llms/bedrock_invoke_agents.py b/litellm/types/llms/bedrock_invoke_agents.py index 62f3c18ca7a..b334c986d6b 100644 --- a/litellm/types/llms/bedrock_invoke_agents.py +++ b/litellm/types/llms/bedrock_invoke_agents.py @@ -4,7 +4,7 @@ Type definitions for AWS Bedrock Invoke Agent API responses. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_InvokeAgent.html """ -from typing import Any, Final +from typing import Final from typing_extensions import TypedDict @@ -39,7 +39,7 @@ class InvokeAgentModelInvocationInput(TypedDict, total=False): """Model invocation input details.""" foundationModel: str | None - inferenceConfiguration: dict[str, Any] | None + inferenceConfiguration: dict[str, object] | None text: str | None traceId: str | None type: str | None @@ -49,9 +49,9 @@ class InvokeAgentModelInvocationOutput(TypedDict, total=False): """Model invocation output details.""" metadata: InvokeAgentMetadata | None - parsedResponse: dict[str, Any] | None - rawResponse: dict[str, Any] | None - reasoningContent: dict[str, Any] | None + parsedResponse: dict[str, object] | None + rawResponse: dict[str, object] | None + reasoningContent: dict[str, object] | None traceId: str | None diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index c300c26bcd3..2ad81e64e05 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -1,5 +1,6 @@ +from collections.abc import Mapping from enum import Enum -from typing import Any, Literal +from typing import Literal from typing_extensions import Required, TypedDict @@ -235,7 +236,7 @@ class GeminiVideoGenerationInstance(TypedDict, total=False): """Instance data for Gemini video generation request""" prompt: Required[str] - image: dict[str, Any] + image: Mapping[str, object] class GeminiVideoGenerationParameters(LiteLLMBaseModel): @@ -263,7 +264,7 @@ class GeminiVideoGenerationParameters(LiteLLMBaseModel): negativePrompt: str | None = None """Text describing what not to include in the video.""" - lastFrame: Any | None = None + lastFrame: object | None = None """ The final image for interpolation video to transition. Should be used with the 'image' parameter. @@ -275,7 +276,7 @@ class GeminiVideoGenerationParameters(LiteLLMBaseModel): Only supported in Veo 3.1 (list of VideoGenerationReferenceImage objects). """ - video: Any | None = None + video: object | None = None """ Video to be used for video extension (Video object). Only supported in Veo 3.1 & Veo 3 Fast. @@ -354,5 +355,5 @@ class GeminiLongRunningOperationResponse(LiteLLMBaseModel): response: GeminiOperationResponse | None = None """Response object when operation is complete""" - error: dict[str, Any] | None = None + error: dict[str, object] | None = None """Error details if operation failed""" diff --git a/litellm/types/llms/langgraph.py b/litellm/types/llms/langgraph.py index 7e34329992d..83f7bb090d1 100644 --- a/litellm/types/llms/langgraph.py +++ b/litellm/types/llms/langgraph.py @@ -4,7 +4,8 @@ Type definitions for LangGraph API. LangGraph provides a streaming and non-streaming API for running agents. """ -from typing import Any, Literal +from collections.abc import Mapping +from typing import Literal from typing_extensions import TypedDict @@ -29,8 +30,8 @@ class LangGraphRequest(TypedDict, total=False): assistant_id: str input: LangGraphInput stream_mode: str | None - config: dict[str, Any] | None - metadata: dict[str, Any] | None + config: Mapping[str, object] | None + metadata: Mapping[str, object] | None # Response Types - Streaming @@ -38,7 +39,7 @@ class LangGraphStreamEvent(TypedDict, total=False): """Single event in a LangGraph stream response.""" event: str - data: Any + data: object # Response Types - Non-streaming @@ -55,7 +56,7 @@ class LangGraphResponse(TypedDict, total=False): """Non-streaming response structure from LangGraph.""" messages: list[LangGraphResponseMessage] - values: dict[str, Any] + values: dict[str, object] # Parsed response for internal use diff --git a/litellm/types/mcp_server/tool_registry.py b/litellm/types/mcp_server/tool_registry.py index 738d37c9503..080e27b6ed3 100644 --- a/litellm/types/mcp_server/tool_registry.py +++ b/litellm/types/mcp_server/tool_registry.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from typing import Any, ClassVar +from typing import ClassVar from pydantic import ConfigDict, Field @@ -10,7 +10,7 @@ class MCPTool(LiteLLMBaseModel): model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True) name: str description: str - input_schema: dict[str, Any] + input_schema: dict[str, object] handler: Callable server_id: str | None = Field(default=None, frozen=True) @@ -18,18 +18,18 @@ class MCPTool(LiteLLMBaseModel): class ToolSchema(LiteLLMBaseModel): name: str description: str - inputSchema: dict[str, Any] + inputSchema: dict[str, object] class ListToolsResponse(LiteLLMBaseModel): tools: list[ToolSchema] nextCursor: str | None = None - _meta: dict[str, Any] | None = None + _meta: dict[str, object] | None = None class CallToolRequest(LiteLLMBaseModel): method: str = "tools/call" - params: dict[str, Any] + params: dict[str, object] class ContentItem(LiteLLMBaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py b/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py index d1eba1dd4a9..f936c9d85cf 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/enkryptai.py @@ -1,5 +1,3 @@ -from typing import Any - from pydantic import Field from typing_extensions import TypedDict @@ -20,7 +18,7 @@ class EnkryptAIPolicyViolationDetail(TypedDict, total=False): class EnkryptAIPIIDetail(TypedDict, total=False): """Details for PII detection.""" - pii: dict[str, Any] + pii: dict[str, object] class EnkryptAIToxicityDetail(TypedDict, total=False): @@ -85,8 +83,8 @@ class EnkryptAIResponseDetails(TypedDict, total=False): toxicity: EnkryptAIToxicityDetail | None keyword_detected: EnkryptAIKeywordDetail | None bias: EnkryptAIBiasDetail | None - prompt_injection: dict[str, Any] | None - jailbreak: dict[str, Any] | None + prompt_injection: dict[str, object] | None + jailbreak: dict[str, object] | None class EnkryptAIResponse(TypedDict, total=False): @@ -107,7 +105,7 @@ class EnkryptAIProcessedResult(TypedDict): | EnkryptAIToxicityDetail | EnkryptAIKeywordDetail | EnkryptAIBiasDetail - | dict[str, Any], + | dict[str, object], ] diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 213c911cfdc..3f1a7980b27 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -65,6 +65,34 @@ class MetricWithMetadata(MetricBase): api_key_breakdown: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} +class GuardrailSpendMetrics(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + spend: float = 0.0 + untracked_units: int = 0 + + +class GuardrailKeyMetric(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + metrics: GuardrailSpendMetrics + metadata: KeyMetadata + + +class GuardrailMetadata(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + guardrail_name: str + + +class GuardrailMetric(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + metrics: GuardrailSpendMetrics + metadata: GuardrailMetadata + api_key_breakdown: Mapping[str, GuardrailKeyMetric] = Field(default_factory=dict) + + class BreakdownMetrics(LiteLLMBaseModel): """Breakdown of spend by different dimensions""" @@ -75,6 +103,7 @@ class BreakdownMetrics(LiteLLMBaseModel): endpoints: dict[str, MetricWithMetadata] = Field(default_factory=dict) # endpoint -> {metrics, metadata} api_keys: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} entities: dict[str, MetricWithMetadata] = Field(default_factory=dict) # entity -> {metrics, metadata} + guardrails: Mapping[str, GuardrailMetric] = Field(default_factory=dict) # guardrail_id -> spend class DailySpendData(LiteLLMBaseModel): diff --git a/litellm/types/repositories/daily_activity.py b/litellm/types/repositories/daily_activity.py index 4111b83a891..b013ce43113 100644 --- a/litellm/types/repositories/daily_activity.py +++ b/litellm/types/repositories/daily_activity.py @@ -168,6 +168,16 @@ class EntityRollupRow(RollupMetricsRow): distinct_api_keys: int | None +@dataclass(frozen=True, slots=True) +class GuardrailSpendRow: + date: str + guardrail_id: str + guardrail_name: str + api_key: str + spend: float + untracked_units: int + + @dataclass(frozen=True, slots=True) class AggregatedRows: grouping_rows: tuple[GroupingSetsRow, ...] diff --git a/litellm/types/router.py b/litellm/types/router.py index ba9b82e2a57..f8de2128a80 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -533,6 +533,8 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): rpm: int | None = None itpm: int | None = None otpm: int | None = None + tpd: int | None = None + rpd: int | None = None timeout: float | str | httpx.Timeout | SDKTimeout | None = None # if str, pass in as os.environ/ # read timeout of a live stream, a faked stream takes the longer of this and timeout. If str, pass in as os.environ/ stream_timeout: float | str | None = None @@ -738,6 +740,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): rpm: int | None itpm: int | None otpm: int | None + tpd: ReadOnly[int | None] + rpd: ReadOnly[int | None] order: int | None weight: int | None max_parallel_requests: int | None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 055c18f2e36..b21bf0e14d4 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -225,6 +225,8 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_xhigh_reasoning_effort: bool | None supports_max_reasoning_effort: bool | None reasoning_effort_levels: ReadOnly[Sequence[str] | None] + reasoning_modes: ReadOnly[Sequence[str] | None] + reasoning_contexts: ReadOnly[Sequence[str] | None] default_reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None] supports_output_config: bool | None supports_image_size: bool | None @@ -4586,6 +4588,7 @@ class LlmProviders(str, Enum): PRISM = "prism" OPENINFER = "openinfer" INFERENCEAPIS = "inferenceapis" + ZEROGPU = "zerogpu" DARKBLOOM = "darkbloom" META = "meta" SAIL = "sail" diff --git a/litellm/utils.py b/litellm/utils.py index 1e095966a43..be1e11dd39f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7000,6 +7000,8 @@ def get_model_info_helper( supports_xhigh_reasoning_effort=_model_info.get("supports_xhigh_reasoning_effort", None), supports_max_reasoning_effort=_model_info.get("supports_max_reasoning_effort", None), reasoning_effort_levels=_model_info.get("reasoning_effort_levels", None), + reasoning_modes=_model_info.get("reasoning_modes", None), + reasoning_contexts=_model_info.get("reasoning_contexts", None), default_reasoning_effort=_model_info.get("default_reasoning_effort", None), bedrock_output_config_effort_ceiling=_model_info.get("bedrock_output_config_effort_ceiling", None), bedrock_converse_supports_strict_tools=_model_info.get("bedrock_converse_supports_strict_tools", None), @@ -7662,7 +7664,6 @@ def validate_environment( model in litellm.open_ai_chat_completion_models or model in litellm.open_ai_text_completion_models or model in litellm.open_ai_embedding_models - or model in litellm.openai_image_generation_models or model.startswith("gpt-image") ): if "OPENAI_API_KEY" in os.environ: @@ -9612,18 +9613,19 @@ class ProviderConfigManager: return litellm.HostedVLLMDecisionsConfig() if provider == LlmProviders.LM_STUDIO: return litellm.LMStudioDecisionsConfig() + if provider in (LlmProviders.OLLAMA, LlmProviders.OLLAMA_CHAT): + return litellm.OllamaDecisionsConfig() if provider == LlmProviders.LAYA: return litellm.LayaDecisionsConfig() if provider == LlmProviders.BESPOKE: return litellm.BespokeDecisionsConfig() - if provider in (LlmProviders.OPENAI, LlmProviders.AZURE): + if provider in (LlmProviders.OPENAI, LlmProviders.AZURE, LlmProviders.AZURE_AI): configs: Final = { LlmProviders.OPENAI: litellm.OpenAIDecisionsConfig, LlmProviders.AZURE: litellm.AzureDecisionsConfig, + LlmProviders.AZURE_AI: litellm.AzureAIDecisionsConfig, } return configs[provider]() - if provider == LlmProviders.AZURE_AI: - return litellm.AzureAIDecisionsConfig() return None @staticmethod @@ -10308,11 +10310,9 @@ class ProviderConfigManager: provider: LlmProviders, ) -> BaseImageGenerationConfig | None: if LlmProviders.OPENAI == provider: - from litellm.llms.openai.image_generation import ( - get_openai_image_generation_config, - ) + from litellm.llms.openai.image_generation import GPTImageGenerationConfig - return get_openai_image_generation_config(model) + return GPTImageGenerationConfig() elif LlmProviders.AZURE == provider: from litellm.llms.azure.image_generation import ( get_azure_image_generation_config, @@ -10415,6 +10415,13 @@ class ProviderConfigManager: return get_modelscope_image_generation_config(model) elif LlmProviders.EDENAI == provider: return litellm.EdenAIImageGenerationConfig() + elif LlmProviders.TENCENT == provider: + from litellm.llms.tencent.image_generation.transformation import ( + TencentImageGenerationConfig, + is_tokenhub_image_model, + ) + + return TencentImageGenerationConfig() if is_tokenhub_image_model(model) else None return None @staticmethod @@ -10536,9 +10543,9 @@ class ProviderConfigManager: provider: LlmProviders, ) -> BaseImageEditConfig | None: if LlmProviders.OPENAI == provider: - from litellm.llms.openai.image_edit import get_openai_image_edit_config + from litellm.llms.openai.image_edit import OpenAIImageEditConfig - return get_openai_image_edit_config(model=model) + return OpenAIImageEditConfig() elif LlmProviders.HOSTED_VLLM == provider: from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config @@ -10603,6 +10610,15 @@ class ProviderConfigManager: ) return get_openrouter_image_edit_config(model) + elif LlmProviders.TENCENT == provider: + from litellm.llms.tencent.image_edit.transformation import ( + TencentImageEditConfig, + ) + from litellm.llms.tencent.image_generation.transformation import ( + is_tokenhub_image_model, + ) + + return TencentImageEditConfig() if is_tokenhub_image_model(model) else None return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1b65c87c315..8719565a96f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -21009,10 +21009,11 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-fable-5": { - "cache_creation_input_token_cost": 1.250004e-05, + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_1hr": 1.999998e-05, "cache_read_input_token_cost": 1.00002e-06, - "input_cost_per_token": 1.000006e-05, - "input_dbu_cost_per_token": 0.000142858, + "input_cost_per_token": 9.99999e-06, + "input_dbu_cost_per_token": 0.00014286, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -21021,7 +21022,7 @@ }, "mode": "chat", "output_cost_per_token": 5.000002e-05, - "output_dbu_cost_per_token": 0.000714286, + "output_dbu_cost_per_token": 0.00071429, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21037,10 +21038,11 @@ "thinking_always_on": true }, "databricks/databricks-claude-fable-5-1": { - "cache_creation_input_token_cost": 1.250004e-05, - "cache_read_input_token_cost": 2.5004e-07, - "input_cost_per_token": 1.000006e-05, - "input_dbu_cost_per_token": 0.000142858, + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_1hr": 1.999998e-05, + "cache_read_input_token_cost": 2.4997e-07, + "input_cost_per_token": 9.99999e-06, + "input_dbu_cost_per_token": 0.00014286, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, @@ -21049,7 +21051,7 @@ }, "mode": "chat", "output_cost_per_token": 5.000002e-05, - "output_dbu_cost_per_token": 0.000714286, + "output_dbu_cost_per_token": 0.00071429, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21067,6 +21069,7 @@ }, "databricks/databricks-claude-haiku-4-5": { "cache_creation_input_token_cost": 1.24999e-06, + "cache_creation_input_token_cost_above_1hr": 1.99997e-06, "cache_read_input_token_cost": 1.0003e-07, "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, @@ -21090,9 +21093,10 @@ }, "databricks/databricks-claude-opus-4": { "cache_creation_input_token_cost": 1.874999e-05, + "cache_creation_input_token_cost_above_1hr": 2.999997e-05, "cache_read_input_token_cost": 1.50003e-06, - "input_cost_per_token": 1.5000020000000002e-05, - "input_dbu_cost_per_token": 0.000214286, + "input_cost_per_token": 1.500002e-05, + "input_dbu_cost_per_token": 0.00021429, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -21100,8 +21104,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 7.500003000000001e-05, - "output_dbu_cost_per_token": 0.001071429, + "output_cost_per_token": 7.50001e-05, + "output_dbu_cost_per_token": 0.0010714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21113,9 +21117,10 @@ }, "databricks/databricks-claude-opus-4-1": { "cache_creation_input_token_cost": 1.874999e-05, + "cache_creation_input_token_cost_above_1hr": 2.999997e-05, "cache_read_input_token_cost": 1.50003e-06, - "input_cost_per_token": 1.5000020000000002e-05, - "input_dbu_cost_per_token": 0.000214286, + "input_cost_per_token": 1.500002e-05, + "input_dbu_cost_per_token": 0.00021429, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 32000, @@ -21123,8 +21128,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 7.500003000000001e-05, - "output_dbu_cost_per_token": 0.001071429, + "output_cost_per_token": 7.50001e-05, + "output_dbu_cost_per_token": 0.0010714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21136,6 +21141,7 @@ }, "databricks/databricks-claude-opus-4-5": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21146,8 +21152,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.00035714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21160,6 +21166,7 @@ }, "databricks/databricks-claude-opus-4-6": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21170,8 +21177,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.5000010000000002e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_cost_per_token": 2.500001e-05, + "output_dbu_cost_per_token": 0.00035714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": false, "supports_function_calling": true, @@ -21184,6 +21191,7 @@ }, "databricks/databricks-claude-opus-4-7": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21195,7 +21203,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 2048, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21210,6 +21218,7 @@ }, "databricks/databricks-claude-opus-4-8": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21221,7 +21230,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21237,6 +21246,7 @@ }, "databricks/databricks-claude-opus-5": { "cache_creation_input_token_cost": 6.25002e-06, + "cache_creation_input_token_cost_above_1hr": 9.99999e-06, "cache_read_input_token_cost": 5.0001e-07, "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, @@ -21248,7 +21258,7 @@ }, "mode": "chat", "output_cost_per_token": 2.500001e-05, - "output_dbu_cost_per_token": 0.000357143, + "output_dbu_cost_per_token": 0.00035714, "prompt_cache_min_tokens": 512, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21344,8 +21354,9 @@ }, "databricks/databricks-claude-sonnet-4-5": { "cache_creation_input_token_cost": 3.74997e-06, + "cache_creation_input_token_cost_above_1hr": 5.99998e-06, "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 200000, @@ -21354,8 +21365,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.00021429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, @@ -21367,8 +21378,9 @@ }, "databricks/databricks-claude-sonnet-4-6": { "cache_creation_input_token_cost": 3.74997e-06, + "cache_creation_input_token_cost_above_1hr": 5.99998e-06, "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.9999900000000002e-06, + "input_cost_per_token": 2.99999e-06, "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, @@ -21377,8 +21389,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5000020000000002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 1.500002e-05, + "output_dbu_cost_per_token": 0.00021429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": false, "supports_function_calling": true, @@ -21390,19 +21402,20 @@ "prompt_cache_min_tokens": 1024 }, "databricks/databricks-claude-sonnet-5": { - "cache_creation_input_token_cost": 3.74997e-06, - "cache_read_input_token_cost": 3.0002e-07, - "input_cost_per_token": 2.99999e-06, - "input_dbu_cost_per_token": 4.2857e-05, + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_1hr": 4.00001e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 1000000, "max_output_tokens": 128000, "metadata": { - "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. Introductory launch rates of 28.571 input / 142.857 output / 35.714 cache write / 2.857 cache read DBU run through 2026-08-31; the standard rates are listed here because entries carry no expiry date." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 1.500002e-05, - "output_dbu_cost_per_token": 0.000214286, + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, "prompt_cache_min_tokens": 1024, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_adaptive_thinking": true, @@ -21546,7 +21559,10 @@ }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_audio_token_cost": 6.251e-08, "cache_read_input_token_cost": 3.122e-08, + "input_cost_per_audio_token": 6.2503e-07, + "input_cost_per_image_token": 3.1248e-07, "input_cost_per_token": 3.1248e-07, "input_dbu_cost_per_token": 4.464e-06, "litellm_provider": "databricks", @@ -21573,16 +21589,28 @@ ], "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_video_input": true, + "input_cost_per_token_priority": 5.6252e-07, + "output_cost_per_token_priority": 3.37498e-06, + "cache_read_input_token_cost_priority": 5.628e-08, + "cache_creation_input_token_cost_priority": 5.6252e-07, + "input_cost_per_audio_token_priority": 1.12497e-06, + "cache_read_input_audio_token_cost_priority": 1.1249e-07 }, "databricks/databricks-gemini-3-1-flash-image": { + "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_token_cost": 6.2503e-07, + "input_cost_per_image_token": 6.2503e-07, + "input_cost_per_token": 6.2503e-07, "litellm_provider": "databricks", "max_input_tokens": 131072, "max_output_tokens": 32768, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_image_token": 7.50001e-05, + "output_cost_per_token": 3.74997e-06, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21596,13 +21624,19 @@ "supports_vision": true }, "databricks/databricks-gemini-3-pro-image": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 2.49998e-06, + "input_cost_per_image_token": 2.49998e-06, + "input_cost_per_token": 2.49998e-06, "litellm_provider": "databricks", "max_input_tokens": 65536, "max_output_tokens": 32768, "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", + "output_cost_per_image_token": 0.0001500002, + "output_cost_per_token": 1.500002e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21618,7 +21652,9 @@ "databricks/databricks-gemini-3-1-pro": { "cache_creation_input_token_cost": 2.49998e-06, "cache_read_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost_above_200k_tokens": 5.0001e-07, "input_cost_per_token": 2.49998e-06, + "input_cost_per_token_above_200k_tokens": 5.00003e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", "max_input_tokens": 1048576, @@ -21628,12 +21664,9 @@ }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token_above_200k_tokens": 2.25e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_function_calling": true, - "supports_prompt_caching": true, - "supports_tool_choice": true, - "supports_reasoning": true, "supported_modalities": [ "text", "image", @@ -21643,13 +21676,27 @@ "supported_output_modalities": [ "text" ], - "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "input_cost_per_token_priority": 4.50002e-06, + "output_cost_per_token_priority": 2.699998e-05, + "cache_read_input_token_cost_priority": 4.5003e-07, + "cache_creation_input_token_cost_priority": 4.50002e-06, + "input_cost_per_token_above_200k_tokens_priority": 8.99997e-06, + "output_cost_per_token_above_200k_tokens_priority": 4.049997e-05, + "cache_read_input_token_cost_above_200k_tokens_priority": 8.9999e-07 }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, + "cache_read_input_audio_token_cost": 1.2502e-07, "cache_read_input_token_cost": 6.251e-08, + "input_cost_per_audio_token": 1.24999e-06, + "input_cost_per_image_token": 6.2503e-07, "input_cost_per_token": 6.2503e-07, "input_dbu_cost_per_token": 8.929e-06, "litellm_provider": "databricks", @@ -21676,7 +21723,13 @@ ], "supports_vision": true, "supports_audio_input": true, - "supports_video_input": true + "supports_video_input": true, + "input_cost_per_token_priority": 1.12497e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.1249e-07, + "cache_creation_input_token_cost_priority": 1.12497e-06, + "input_cost_per_audio_token_priority": 2.25001e-06, + "cache_read_input_audio_token_cost_priority": 2.2498e-07 }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21698,58 +21751,9 @@ "supports_tool_choice": true }, "databricks/databricks-gemini-3-8-flash": { - "litellm_provider": "databricks", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." - }, - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_function_calling": true, - "supports_reasoning": true, - "supports_sampling_params": false, - "supports_tool_choice": true, - "supports_vision": true - }, - "databricks/databricks-gemini-3-7-flash": { - "litellm_provider": "databricks", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "metadata": { - "notes": "Databricks DBU rates not yet published for this model; endpoint metadata only." - }, - "mode": "chat", - "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text" - ], - "supports_function_calling": true, - "supports_reasoning": true, - "supports_sampling_params": false, - "supports_tool_choice": true, - "supports_vision": true - }, - "databricks/databricks-gemini-3-6-flash": { - "cache_creation_input_token_cost": 1.87502e-06, - "cache_read_input_token_cost": 1.8753e-07, - "input_cost_per_token": 1.87502e-06, - "input_dbu_cost_per_token": 2.6786e-05, + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, "litellm_provider": "databricks", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -21757,8 +21761,73 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "chat", - "output_cost_per_token": 9.37503e-06, - "output_dbu_cost_per_token": 0.000133929, + "output_cost_per_token": 3.74997e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 + }, + "databricks/databricks-gemini-3-7-flash": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 3.74997e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 + }, + "databricks/databricks-gemini-3-6-flash": { + "cache_creation_input_token_cost": 7.4998e-07, + "cache_read_input_token_cost": 7.497e-08, + "input_cost_per_token": 7.4998e-07, + "input_dbu_cost_per_token": 1.0714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 3.74997e-06, + "output_dbu_cost_per_token": 5.3571e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -21774,7 +21843,11 @@ "supports_reasoning": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.35002e-06, + "output_cost_per_token_priority": 6.75003e-06, + "cache_read_input_token_cost_priority": 1.3503e-07, + "cache_creation_input_token_cost_priority": 1.35002e-06 }, "databricks/databricks-gemini-3-5-flash": { "cache_creation_input_token_cost": 1.87502e-06, @@ -21804,7 +21877,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 3.37498e-06, + "output_cost_per_token_priority": 2.025002e-05, + "cache_read_input_token_cost_priority": 3.3747e-07, + "cache_creation_input_token_cost_priority": 3.37498e-06 }, "databricks/databricks-gemini-3-5-flash-lite": { "cache_creation_input_token_cost": 3.7499e-07, @@ -21836,7 +21913,11 @@ "supports_response_schema": true, "supports_sampling_params": false, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 6.7501e-07, + "output_cost_per_token_priority": 5.62499e-06, + "cache_read_input_token_cost_priority": 6.748e-08, + "cache_creation_input_token_cost_priority": 6.7501e-07 }, "databricks/databricks-gemma-3-12b": { "cache_creation_input_token_cost": 1.5001e-07, @@ -21879,7 +21960,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": false + "supports_vision": false, + "input_cost_per_token_priority": 2.45e-06, + "output_cost_per_token_priority": 7.7e-06, + "cache_read_input_token_cost_priority": 4.55e-07, + "cache_creation_input_token_cost_priority": 2.45e-06 }, "databricks/databricks-glm-5-3": { "cache_creation_input_token_cost": 1.4e-06, @@ -21909,7 +21994,11 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 2.45e-06, + "output_cost_per_token_priority": 7.7e-06, + "cache_read_input_token_cost_priority": 4.55e-07, + "cache_creation_input_token_cost_priority": 2.45e-06 }, "databricks/databricks-glm-5-3-flash": { "cache_creation_input_token_cost": 1.5001e-07, @@ -21938,7 +22027,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": true, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 2.625e-07, + "output_cost_per_token_priority": 8.75e-07, + "cache_read_input_token_cost_priority": 5.25e-08, + "cache_creation_input_token_cost_priority": 2.625e-07 }, "databricks/databricks-gpt-5": { "cache_creation_input_token_cost": 1.24999e-06, @@ -21957,7 +22050,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 2.49998e-06, + "output_cost_per_token_priority": 1.999998e-05, + "cache_read_input_token_cost_priority": 2.4997e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06 }, "databricks/databricks-gpt-5-1": { "cache_creation_input_token_cost": 1.24999e-06, @@ -21976,7 +22073,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 2.49998e-06, + "output_cost_per_token_priority": 1.999998e-05, + "cache_read_input_token_cost_priority": 2.4997e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06 }, "databricks/databricks-gpt-5-2": { "cache_creation_input_token_cost": 1.75e-06, @@ -21995,7 +22096,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 3.5e-06, + "output_cost_per_token_priority": 2.8e-05, + "cache_read_input_token_cost_priority": 3.5e-07, + "cache_creation_input_token_cost_priority": 3.5e-06 }, "databricks/databricks-gpt-5-3-codex": { "cache_creation_input_token_cost": 1.75e-06, @@ -22014,12 +22119,18 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 3.5e-06, + "output_cost_per_token_priority": 2.8e-05, + "cache_read_input_token_cost_priority": 3.5e-07, + "cache_creation_input_token_cost_priority": 3.5e-06 }, "databricks/databricks-gpt-5-4": { "cache_creation_input_token_cost": 2.49998e-06, "cache_read_input_token_cost": 2.4997e-07, + "cache_read_input_token_cost_above_272k_tokens": 5.0001e-07, "input_cost_per_token": 2.49998e-06, + "input_cost_per_token_above_272k_tokens": 5.00003e-06, "input_dbu_cost_per_token": 3.5714e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22029,6 +22140,7 @@ }, "mode": "chat", "output_cost_per_token": 1.5000020000000002e-05, + "output_cost_per_token_above_272k_tokens": 2.25e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22042,7 +22154,14 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.00003e-06, + "output_cost_per_token_priority": 2.999997e-05, + "cache_read_input_token_cost_priority": 5.0001e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 5.00003e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.999997e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 5.0001e-07 }, "databricks/databricks-gpt-5-4-mini": { "cache_creation_input_token_cost": 7.4998e-07, @@ -22070,7 +22189,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.50003e-06, + "output_cost_per_token_priority": 8.99997e-06, + "cache_read_input_token_cost_priority": 1.5001e-07, + "cache_creation_input_token_cost_priority": 1.50003e-06 }, "databricks/databricks-gpt-5-4-nano": { "cache_creation_input_token_cost": 1.9999e-07, @@ -22102,8 +22225,11 @@ }, "databricks/databricks-gpt-5-6-sol": { "cache_creation_input_token_cost": 5.00003e-06, + "cache_creation_input_token_cost_above_272k_tokens": 9.99999e-06, "cache_read_input_token_cost": 3.9998e-07, + "cache_read_input_token_cost_above_272k_tokens": 8.0003e-07, "input_cost_per_token": 4.00001e-06, + "input_cost_per_token_above_272k_tokens": 8.00002e-06, "input_dbu_cost_per_token": 5.7143e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22113,6 +22239,7 @@ }, "mode": "chat", "output_cost_per_token": 1.999998e-05, + "output_cost_per_token_above_272k_tokens": 3e-05, "output_dbu_cost_per_token": 0.000285714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22126,12 +22253,23 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 8.00002e-06, + "output_cost_per_token_priority": 4.000003e-05, + "cache_read_input_token_cost_priority": 8.0003e-07, + "cache_creation_input_token_cost_priority": 9.99999e-06, + "input_cost_per_token_above_272k_tokens_priority": 1.599997e-05, + "output_cost_per_token_above_272k_tokens_priority": 6.000001e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.59999e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.999998e-05 }, "databricks/databricks-gpt-5-6-terra": { "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, "cache_read_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost_above_272k_tokens": 3.9998e-07, "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, "input_dbu_cost_per_token": 2.8571e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22141,6 +22279,7 @@ }, "mode": "chat", "output_cost_per_token": 1.200003e-05, + "output_cost_per_token_above_272k_tokens": 1.8e-05, "output_dbu_cost_per_token": 0.000171429, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22154,12 +22293,23 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 2.399999e-05, + "cache_read_input_token_cost_priority": 3.9998e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 3.600002e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 }, "databricks/databricks-gpt-5-6-luna": { "cache_creation_input_token_cost": 2.4997e-07, + "cache_creation_input_token_cost_above_272k_tokens": 5.0001e-07, "cache_read_input_token_cost": 2.002e-08, + "cache_read_input_token_cost_above_272k_tokens": 3.997e-08, "input_cost_per_token": 1.9999e-07, + "input_cost_per_token_above_272k_tokens": 3.9998e-07, "input_dbu_cost_per_token": 2.857e-06, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22169,6 +22319,7 @@ }, "mode": "chat", "output_cost_per_token": 1.20001e-06, + "output_cost_per_token_above_272k_tokens": 1.79998e-06, "output_dbu_cost_per_token": 1.7143e-05, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22182,12 +22333,22 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 3.9998e-07, + "output_cost_per_token_priority": 2.40002e-06, + "cache_read_input_token_cost_priority": 3.997e-08, + "cache_creation_input_token_cost_priority": 5.0001e-07, + "input_cost_per_token_above_272k_tokens_priority": 8.0003e-07, + "output_cost_per_token_above_272k_tokens_priority": 3.60003e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.001e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 1.00002e-06 }, "databricks/databricks-gpt-5-5": { "cache_creation_input_token_cost": 5.00003e-06, "cache_read_input_token_cost": 5.0001e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.00002e-06, "input_cost_per_token": 5.00003e-06, + "input_cost_per_token_above_272k_tokens": 9.99999e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", "max_input_tokens": 922000, @@ -22197,6 +22358,7 @@ }, "mode": "responses", "output_cost_per_token": 2.999997e-05, + "output_cost_per_token_above_272k_tokens": 4.5e-05, "output_dbu_cost_per_token": 0.000428571, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ @@ -22210,13 +22372,21 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 1.249997e-05, + "output_cost_per_token_priority": 7.50001e-05, + "cache_read_input_token_cost_priority": 1.24999e-06, + "cache_creation_input_token_cost_priority": 1.249997e-05, + "input_cost_per_token_above_272k_tokens_priority": 1.249997e-05, + "output_cost_per_token_above_272k_tokens_priority": 7.50001e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.24999e-06 }, "databricks/databricks-gpt-5-5-pro": { "cache_creation_input_token_cost": 2.999997e-05, "cache_read_input_token_cost": 2.999997e-05, "input_cost_per_token": 2.999997e-05, - "input_dbu_cost_per_token": 0.000428571, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_dbu_cost_per_token": 0.00042857, "litellm_provider": "databricks", "max_input_tokens": 922000, "max_output_tokens": 128000, @@ -22224,8 +22394,9 @@ "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." }, "mode": "responses", - "output_cost_per_token": 0.00018000003, - "output_dbu_cost_per_token": 0.002571429, + "output_cost_per_token": 0.0001800001, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_dbu_cost_per_token": 0.0025714, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supported_modalities": [ "text", @@ -22257,7 +22428,11 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "input_cost_per_token_priority": 4.5003e-07, + "output_cost_per_token_priority": 3.60003e-06, + "cache_read_input_token_cost_priority": 4.501e-08, + "cache_creation_input_token_cost_priority": 4.5003e-07 }, "databricks/databricks-gpt-5-nano": { "cache_creation_input_token_cost": 4.998e-08, @@ -22407,7 +22582,11 @@ "supports_prompt_caching": true, "supports_reasoning": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "input_cost_per_token_priority": 5.25e-06, + "output_cost_per_token_priority": 2.625e-05, + "cache_read_input_token_cost_priority": 5.25e-07, + "cache_creation_input_token_cost_priority": 5.25e-06 }, "databricks/system.ai.kimi-k3": { "litellm_provider": "databricks", @@ -22489,7 +22668,11 @@ "supports_reasoning": true, "supports_tool_choice": true, "supports_vision": false, - "thinking_always_on": true + "thinking_always_on": true, + "input_cost_per_token_priority": 3.85e-07, + "output_cost_per_token_priority": 3.85e-06, + "cache_read_input_token_cost_priority": 3.85e-07, + "cache_creation_input_token_cost_priority": 3.85e-07 }, "databricks/databricks-qwen3-next-80b-a3b-instruct": { "cache_creation_input_token_cost": 1.5001e-07, @@ -22544,6 +22727,25 @@ "output_dbu_cost_per_token": 0.0, "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models" }, + "databricks/databricks-typesafeai-jev-1-14-0": { + "cache_creation_input_token_cost": 4.2e-08, + "cache_read_input_token_cost": 4.2e-08, + "input_cost_per_token": 4.2e-08, + "input_dbu_cost_per_token": 6e-07, + "litellm_provider": "databricks", + "metadata": { + "notes": "TypeSafe AI Jev 1.14.0, provider-hosted by TypeSafe AI and served through Unity Gateway, queried through the System One API. The proprietary pricing page lists 'Jev 1.14' at 0.600 DBU per 1M input tokens; input_cost_per_token is that times $0.070 per DBU. Output is free. No cache rates are published, so cache tokens bill at the input rate." + }, + "mode": "evaluation", + "output_cost_per_token": 0.0, + "output_dbu_cost_per_token": 0.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supports_vision": true + }, "databricks/ai_decide": { "cache_creation_input_token_cost": 0.0, "cache_read_input_token_cost": 0.0, @@ -28651,6 +28853,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -28708,6 +28911,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -28761,12 +28965,13 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -28812,6 +29017,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -28857,9 +29063,10 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -28904,6 +29111,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -29470,10 +29678,12 @@ }, "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, "google_maps_grounding_cost_per_query": 0.025, + "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, "input_cost_per_token_priority": 2.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, @@ -29488,6 +29698,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -29499,6 +29710,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -30027,6 +30239,7 @@ "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_batches": 2e-06, "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, "input_cost_per_token_flex": 1e-06, @@ -30038,6 +30251,7 @@ "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "output_cost_per_token_flex": 6e-06, @@ -35598,6 +35812,13 @@ "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35675,6 +35896,13 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35851,6 +36079,13 @@ "mode": "chat", "output_cost_per_token": 1e-05, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -35902,6 +36137,13 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36017,6 +36259,13 @@ "mode": "chat", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36066,6 +36315,13 @@ "mode": "responses", "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36107,6 +36363,13 @@ "mode": "responses", "output_cost_per_token": 0.000168, "output_cost_per_token_batches": 8.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -36194,6 +36457,15 @@ "output_cost_per_token_flex": 2.5e-05, "output_cost_per_token_priority": 0.0001, "output_cost_per_token_ultrafast": 0.0003, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36278,6 +36550,15 @@ "output_cost_per_token_above_272k_tokens_batches": 7.5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36370,6 +36651,15 @@ "output_cost_per_token_above_272k_tokens_batches": 3.75e-07, "output_cost_per_token_flex": 2.5e-07, "output_cost_per_token_priority": 1e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36453,6 +36743,15 @@ "output_cost_per_token_above_272k_tokens_batches": 1.5e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36534,6 +36833,15 @@ "output_cost_per_token_above_272k_tokens_batches": 1.5e-05, "output_cost_per_token_flex": 1e-05, "output_cost_per_token_priority": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36617,6 +36925,15 @@ "output_cost_per_token_above_272k_tokens_batches": 9e-06, "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36699,6 +37016,15 @@ "output_cost_per_token_above_272k_tokens_batches": 9e-07, "output_cost_per_token_flex": 6e-07, "output_cost_per_token_priority": 2.4e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -36936,6 +37262,14 @@ "reasoning_effort_levels": [ "medium" ], + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "source": "https://developers.openai.com/api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -36983,6 +37317,14 @@ "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37045,6 +37387,14 @@ "output_cost_per_token_batches": 1.5e-05, "output_cost_per_token_above_272k_tokens_batches": 2.25e-05, "output_cost_per_token_priority": 7.5e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37099,6 +37449,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37150,6 +37508,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37209,6 +37575,14 @@ "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_priority": 3e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -37266,6 +37640,14 @@ "output_cost_per_token_batches": 7.5e-06, "output_cost_per_token_above_272k_tokens_batches": 1.125e-05, "output_cost_per_token_priority": 3e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "supported_endpoints": [ @@ -37315,6 +37697,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37367,6 +37757,14 @@ "output_cost_per_token_flex": 9e-05, "output_cost_per_token_batches": 9e-05, "output_cost_per_token_above_272k_tokens_batches": 0.000135, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37421,6 +37819,14 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37475,6 +37881,14 @@ "output_cost_per_token_flex": 2.25e-06, "output_cost_per_token_batches": 2.25e-06, "output_cost_per_token_priority": 9e-06, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37526,6 +37940,14 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37578,6 +38000,14 @@ "output_cost_per_token": 1.25e-06, "output_cost_per_token_flex": 6.25e-07, "output_cost_per_token_batches": 6.25e-07, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -37624,6 +38054,13 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37666,6 +38103,13 @@ "mode": "responses", "output_cost_per_token": 0.00012, "output_cost_per_token_batches": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37716,6 +38160,13 @@ "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37826,6 +38277,13 @@ "mode": "responses", "output_cost_per_token": 1.4e-05, "output_cost_per_token_priority": 2.8e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37908,6 +38366,13 @@ "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -37959,6 +38424,13 @@ "output_cost_per_token_batches": 1e-06, "output_cost_per_token_flex": 1e-06, "output_cost_per_token_priority": 3.6e-06, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38007,6 +38479,13 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -38055,6 +38534,13 @@ "output_cost_per_token": 4e-07, "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -41845,6 +42331,22 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "mistral/deepseek-v4-1-flash-1017": { + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "mistral", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.mistral.ai/models/deepseek-v4-flash-1017", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "mistral/zai-glm-5-3": { "cache_read_input_token_cost": 1.4e-07, "input_cost_per_token": 1.4e-06, @@ -43805,6 +44307,13 @@ "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -43913,6 +44422,13 @@ "output_cost_per_token_batches": 4e-06, "output_cost_per_token_flex": 4e-06, "output_cost_per_token_priority": 1.4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "standard" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -43992,6 +44508,13 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -44029,6 +44552,13 @@ "mode": "responses", "output_cost_per_token": 8e-05, "output_cost_per_token_batches": 4e-05, + "reasoning_contexts": [ + "auto", + "current_turn" + ], + "reasoning_modes": [ + "pro" + ], "search_context_cost_per_query": { "search_context_size_high": 0.01, "search_context_size_low": 0.01, @@ -51392,6 +51922,7 @@ "mode": "image_generation", "output_cost_per_image": 0.039, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_reasoning_token": 2.5e-06, "output_cost_per_token": 2.5e-06, "output_cost_per_token_batches": 1.25e-06, @@ -51450,6 +51981,7 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, @@ -51484,13 +52016,14 @@ "mode": "image_generation", "output_cost_per_image": 0.134, "output_cost_per_image_token": 0.00012, + "output_cost_per_image_token_batches": 6e-05, "output_cost_per_token": 1.2e-05, "output_cost_per_token_priority": 2.16e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, "supports_reasoning": false, - "source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -51516,6 +52049,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, @@ -51543,10 +52077,11 @@ "mode": "image_generation", "output_cost_per_image": 0.0672, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "supports_reasoning": false, - "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models", + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/images/generations", @@ -51571,6 +52106,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 1.5e-06, "output_cost_per_token_batches": 7.5e-07, "output_cost_per_token_flex": 7.5e-07, @@ -61468,6 +62004,39 @@ "supports_reasoning": true, "supports_vision": false }, + "tencent/seedream-image-v5.0-flash": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.018, + "output_cost_per_image_token": 1.6e-06, + "source": "https://cloud.tencent.com/document/product/1823/130055", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "tencent/seedream-image-v5.0-lite": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.035, + "output_cost_per_image_token": 1.6e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, + "tencent/seedream-image-v5.0-pro": { + "litellm_provider": "tencent", + "mode": "image_generation", + "output_cost_per_image": 0.09, + "output_cost_per_image_token": 1.6e-06, + "source": "https://www.tencentcloud.com/document/product/1300/78937", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ] + }, "cognition/swe-1.6": { "input_cost_per_token": 5e-07, "output_cost_per_token": 2.5e-06, @@ -65847,6 +66416,180 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "zerogpu/deepseek-v4.1-flash": { + "cache_read_input_token_cost": 3e-09, + "input_cost_per_token": 3e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/glm-5.3-flash": { + "cache_read_input_token_cost": 1.5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1048576, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-4.1-mini": { + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 4e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.6e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-5.4-nano": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.25e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-5.6-luna": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/gpt-6-luna": { + "cache_read_input_token_cost": 1e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 272000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "zerogpu/gpt-oss-120b": { + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 128000, + "max_output_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/LFM2.5-1.2B-Instruct": { + "input_cost_per_token": 2e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-08, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/LFM2.5-1.2B-Thinking": { + "input_cost_per_token": 2e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 5e-08, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/llama-3.1-8b-instruct-fast": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/llama-3.2-3b-instruct-fast": { + "cache_read_input_token_cost": 2.5e-08, + "input_cost_per_token": 5e-08, + "litellm_provider": "zerogpu", + "max_input_tokens": 131072, + "max_output_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 3e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog" + }, + "zerogpu/qwen3-30b-a3b-fp8": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 1e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 32768, + "max_output_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 4.5e-07, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, + "zerogpu/qwen3.6-35b-a3b": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 2e-07, + "litellm_provider": "zerogpu", + "max_input_tokens": 262144, + "max_output_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.5e-06, + "source": "https://docs.zerogpu.ai/docs/model-catalog", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true + }, "together_ai/Qwen/Qwen3.8-Flash": { "input_cost_per_token": 1.5e-07, "litellm_provider": "together_ai", @@ -73272,6 +74015,15 @@ "output_cost_per_token_flex": 5e-06, "output_cost_per_token_priority": 2e-05, "output_cost_per_token_ultrafast": 6e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": { @@ -73666,6 +74418,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, @@ -73711,6 +74464,7 @@ "mode": "image_generation", "output_cost_per_image": 0.0336, "output_cost_per_image_token": 3e-05, + "output_cost_per_image_token_batches": 1.5e-05, "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, @@ -76771,6 +77525,15 @@ "output_cost_per_token_batches": 2.5e-06, "output_cost_per_token_flex": 2.5e-06, "output_cost_per_token_priority": 1e-05, + "reasoning_contexts": [ + "auto", + "current_turn", + "all_turns" + ], + "reasoning_modes": [ + "standard", + "pro" + ], "regional_processing_uplift_multiplier_eu": 1.1, "regional_processing_uplift_multiplier_us": 1.1, "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, @@ -77095,6 +77858,577 @@ "supports_xhigh_reasoning_effort": true, "use_openai_responses_path": true }, + "databricks/databricks-claude-haiku-5-5": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_creation_input_token_cost_above_100k_tokens": 6.2503e-07, + "cache_creation_input_token_cost_above_1hr": 1.9999e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.00002e-06, + "cache_read_input_token_cost": 1.001e-08, + "cache_read_input_token_cost_above_100k_tokens": 4.998e-08, + "input_cost_per_token": 1.0003e-07, + "input_cost_per_token_above_100k_tokens": 5.0001e-07, + "input_dbu_cost_per_token": 1.429e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_cost_per_token_above_100k_tokens": 2.49998e-06, + "output_dbu_cost_per_token": 7.143e-06, + "prompt_cache_min_tokens": 512.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_anthropic_thinking_payload": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "thinking_always_on": false + }, + "databricks/databricks-claude-sonnet-5-5": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_1hr": 4.00001e-06, + "cache_read_input_token_cost": 1.0003e-07, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, + "prompt_cache_min_tokens": 512.0, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_adaptive_thinking": true, + "supports_anthropic_thinking_payload": true, + "supports_assistant_prefill": false, + "supports_function_calling": true, + "supports_mid_conversation_system": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": false, + "thinking_always_on": true + }, + "databricks/databricks-deepseek-v4-1-flash": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 3.003e-08, + "deprecation_date": "2026-11-19", + "input_cost_per_token": 3.0002e-07, + "input_dbu_cost_per_token": 4.286e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields. The docs list only databricks-deepseek-v4-1-flash (retiring 2026-11-19) and databricks-deepseek-v4-1-flash-1017, so the pricing page's 'DeepSeek V4.1 Flash (0910)' row prices the unsuffixed endpoint." + }, + "mode": "chat", + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, + "source": "https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 5.25e-07, + "output_cost_per_token_priority": 2.1e-06, + "cache_read_input_token_cost_priority": 5.25e-08, + "cache_creation_input_token_cost_priority": 5.25e-07 + }, + "databricks/databricks-deepseek-v4-1-flash-1017": { + "cache_creation_input_token_cost": 3.0002e-07, + "cache_read_input_token_cost": 7e-09, + "input_cost_per_token": 3.0002e-07, + "input_dbu_cost_per_token": 4.286e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1000000, + "max_output_tokens": 384000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 1.20001e-06, + "output_dbu_cost_per_token": 1.7143e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gemini-2-5-flash-lite": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_read_input_audio_token_cost": 3.752e-08, + "cache_read_input_token_cost": 1.253e-08, + "input_cost_per_audio_token": 3.7499e-07, + "input_cost_per_image_token": 1.2502e-07, + "input_cost_per_token": 1.2502e-07, + "input_dbu_cost_per_token": 1.786e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "databricks/databricks-gemini-3-1-flash-lite-image": { + "cache_creation_input_token_cost": 3.1248e-07, + "cache_read_input_token_cost": 3.1248e-07, + "input_cost_per_image_token": 3.1248e-07, + "input_cost_per_token": 3.1248e-07, + "input_dbu_cost_per_token": 4.464e-06, + "litellm_provider": "databricks", + "max_input_tokens": 65536, + "max_output_tokens": 4096, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_image_token": 3.749998e-05, + "output_cost_per_token": 1.87502e-06, + "output_dbu_cost_per_token": 2.6786e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_vision": true + }, + "databricks/databricks-gemini-4-argon": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 1.2502e-07, + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 1.249997e-05, + "output_dbu_cost_per_token": 0.00017857, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true, + "input_cost_per_token_priority": 4.50002e-06, + "output_cost_per_token_priority": 2.250003e-05, + "cache_read_input_token_cost_priority": 2.2498e-07, + "cache_creation_input_token_cost_priority": 4.50002e-06 + }, + "databricks/databricks-gpt-5-4-pro": { + "cache_creation_input_token_cost": 2.999997e-05, + "cache_read_input_token_cost": 2.999997e-05, + "input_cost_per_token": 2.999997e-05, + "input_cost_per_token_above_272k_tokens": 6e-05, + "input_dbu_cost_per_token": 0.00042857, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 0.0001800001, + "output_cost_per_token_above_272k_tokens": 0.00027, + "output_dbu_cost_per_token": 0.0025714, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-gpt-6-1-sol": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, + "cache_read_input_token_cost": 1.0003e-07, + "cache_read_input_token_cost_above_272k_tokens": 1.9999e-07, + "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 9.99999e-06, + "output_cost_per_token_above_272k_tokens": 1.5e-05, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 1.99997e-05, + "cache_read_input_token_cost_priority": 1.9999e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.99999e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 3.9998e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 + }, + "databricks/databricks-gpt-6-1-terra": { + "cache_creation_input_token_cost": 1.24999e-06, + "cache_creation_input_token_cost_above_272k_tokens": 2.49998e-06, + "cache_read_input_token_cost": 4.998e-08, + "cache_read_input_token_cost_above_272k_tokens": 1.0003e-07, + "input_cost_per_token": 1.00002e-06, + "input_cost_per_token_above_272k_tokens": 1.99997e-06, + "input_dbu_cost_per_token": 1.4286e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.00003e-06, + "output_cost_per_token_above_272k_tokens": 7.50001e-06, + "output_dbu_cost_per_token": 7.1429e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.99997e-06, + "output_cost_per_token_priority": 9.99999e-06, + "cache_read_input_token_cost_priority": 1.0003e-07, + "cache_creation_input_token_cost_priority": 2.49998e-06, + "input_cost_per_token_above_272k_tokens_priority": 4.00001e-06, + "output_cost_per_token_above_272k_tokens_priority": 1.500002e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 1.9999e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.00003e-06 + }, + "databricks/databricks-gpt-6-astra": { + "cache_creation_input_token_cost": 1.249997e-05, + "cache_creation_input_token_cost_above_272k_tokens": 2.5e-05, + "cache_read_input_token_cost": 1.00002e-06, + "cache_read_input_token_cost_above_272k_tokens": 1.99997e-06, + "input_cost_per_token": 9.99999e-06, + "input_cost_per_token_above_272k_tokens": 2e-05, + "input_dbu_cost_per_token": 0.00014286, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 5.000002e-05, + "output_cost_per_token_above_272k_tokens": 7.50001e-05, + "output_dbu_cost_per_token": 0.00071429, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.999998e-05, + "output_cost_per_token_priority": 9.99999e-05, + "cache_read_input_token_cost_priority": 1.99997e-06, + "cache_creation_input_token_cost_priority": 2.500001e-05, + "input_cost_per_token_above_272k_tokens_priority": 4.000003e-05, + "output_cost_per_token_above_272k_tokens_priority": 0.0001500002, + "cache_read_input_token_cost_above_272k_tokens_priority": 4.00001e-06, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.000002e-05 + }, + "databricks/databricks-gpt-6-luna": { + "cache_creation_input_token_cost": 1.2502e-07, + "cache_creation_input_token_cost_above_272k_tokens": 2.4997e-07, + "cache_read_input_token_cost": 1.001e-08, + "cache_read_input_token_cost_above_272k_tokens": 2.002e-08, + "input_cost_per_token": 1.0003e-07, + "input_cost_per_token_above_272k_tokens": 1.9999e-07, + "input_dbu_cost_per_token": 1.429e-06, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 5.0001e-07, + "output_cost_per_token_above_272k_tokens": 7.497e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 1.9999e-07, + "output_cost_per_token_priority": 1.0003e-06, + "cache_read_input_token_cost_priority": 2.002e-08, + "cache_creation_input_token_cost_priority": 2.4997e-07, + "input_cost_per_token_above_272k_tokens_priority": 3.9998e-07, + "output_cost_per_token_above_272k_tokens_priority": 1.5001e-06, + "cache_read_input_token_cost_above_272k_tokens_priority": 3.997e-08, + "cache_creation_input_token_cost_above_272k_tokens_priority": 5.0001e-07 + }, + "databricks/databricks-gpt-6-sol": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_creation_input_token_cost_above_272k_tokens": 5.00003e-06, + "cache_read_input_token_cost": 1.9999e-07, + "cache_read_input_token_cost_above_272k_tokens": 3.9998e-07, + "input_cost_per_token": 1.99997e-06, + "input_cost_per_token_above_272k_tokens": 4.00001e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "responses", + "output_cost_per_token": 9.99999e-06, + "output_cost_per_token_above_272k_tokens": 1.50003e-05, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true, + "input_cost_per_token_priority": 4.00001e-06, + "output_cost_per_token_priority": 1.99997e-05, + "cache_read_input_token_cost_priority": 3.9998e-07, + "cache_creation_input_token_cost_priority": 5.00003e-06, + "input_cost_per_token_above_272k_tokens_priority": 8.00002e-06, + "output_cost_per_token_above_272k_tokens_priority": 2.99999e-05, + "cache_read_input_token_cost_above_272k_tokens_priority": 8.0003e-07, + "cache_creation_input_token_cost_above_272k_tokens_priority": 9.99999e-06 + }, + "databricks/databricks-gpt-image-1": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 9.99999e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 4.000003e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-image-1-5": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 8.00002e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 3.200001e-05, + "output_cost_per_token": 9.99999e-06, + "output_dbu_cost_per_token": 0.00014286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_pdf_input": true, + "supports_vision": true + }, + "databricks/databricks-gpt-image-1-mini": { + "cache_creation_input_token_cost": 1.99997e-06, + "cache_read_input_token_cost": 1.9999e-07, + "input_cost_per_image_token": 2.49998e-06, + "input_cost_per_token": 1.99997e-06, + "input_dbu_cost_per_token": 2.8571e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 8.00002e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-image-2": { + "cache_creation_input_token_cost": 5.00003e-06, + "cache_read_input_token_cost": 1.24999e-06, + "input_cost_per_image_token": 8.00002e-06, + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "image_generation", + "output_cost_per_image_token": 2.999997e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_pdf_input": true, + "supports_vision": true + }, + "databricks/databricks-grok-4-7": { + "cache_creation_input_token_cost": 2.49998e-06, + "cache_read_input_token_cost": 6.2503e-07, + "input_cost_per_token": 2.49998e-06, + "input_dbu_cost_per_token": 3.5714e-05, + "litellm_provider": "databricks", + "max_input_tokens": 500000, + "max_output_tokens": 500000, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 7.50001e-06, + "output_dbu_cost_per_token": 0.00010714, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, + "databricks/databricks-kimi-k2-7": { + "cache_creation_input_token_cost": 9.4997e-07, + "cache_read_input_token_cost": 1.8998e-07, + "input_cost_per_token": 9.4997e-07, + "input_dbu_cost_per_token": 1.3571e-05, + "litellm_provider": "databricks", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "metadata": { + "notes": "Costs per token are the published Global DBU rates times $0.070 per DBU. The '*_dbu_cost_per_token' fields are provided for reference; cost calculation reads the dollar '*_cost_per_token' fields." + }, + "mode": "chat", + "output_cost_per_token": 4.00001e-06, + "output_dbu_cost_per_token": 5.7143e-05, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true + }, "fireworks_ai/accounts/fireworks/routers/spot-deepseek-v4-flash": { "cache_read_input_token_cost": 2.1e-08, "deprecation_date": "2026-08-14", diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 1ab03fb3fdc..0a6c2bcb2a9 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1225,7 +1225,7 @@ "responses": true, "text_completion": true, "embeddings": true, - "image_generations": false, + "image_generations": true, "image_edits": false, "audio_transcriptions": false, "audio_speech": false, @@ -3793,8 +3793,8 @@ "responses": true, "text_completion": true, "embeddings": false, - "image_generations": false, - "image_edits": false, + "image_generations": true, + "image_edits": true, "audio_transcriptions": false, "audio_speech": false, "rerank": false, @@ -4437,6 +4437,32 @@ "generateContent": true } }, + "zerogpu": { + "display_name": "ZeroGPU (`zerogpu`)", + "url": "https://docs.litellm.ai/docs/providers/zerogpu", + "endpoints": { + "chat_completions": true, + "responses": true, + "text_completion": true, + "embeddings": false, + "image_generations": false, + "image_edits": false, + "audio_transcriptions": false, + "audio_speech": false, + "rerank": false, + "ocr": false, + "search": false, + "moderations": false, + "realtime": false, + "video_generations": false, + "batches": false, + "messages": true, + "interactions": true, + "systemone": false, + "a2a": true, + "generateContent": true + } + }, "ragflow": { "display_name": "RAGFlow (`ragflow`)", "url": "https://docs.litellm.ai/docs/providers/ragflow", diff --git a/pyproject.toml b/pyproject.toml index 2d25c1ec737..00cdb1587aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -401,7 +401,7 @@ markers = [ "asyncio: mark test as an asyncio test", "limit_leaks: mark test with memory limit for leak detection (e.g., '40 MB')", "no_parallel: mark test to run sequentially (not in parallel) - typically for memory measurement tests", - "requires_rust_extension: public Python contract requiring an enabled, compiled Rust extension", + "requires_rust_extension: public Python contract tests that run against the release-built wheel", ] filterwarnings = [ # Suppress Pydantic serializer warnings from mock server responses (non-critical for memory tests) diff --git a/schema.prisma b/schema.prisma index 83a0e6f9e05..2965174114b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1163,6 +1163,7 @@ model LiteLLM_ManagedFileTable { @@index([unified_file_id]) @@index([flat_model_file_ids], type: Gin) @@index([team_id, created_at(sort: Desc)]) + @@index([created_by, created_at(sort: Desc)]) } model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use the @@ -2060,3 +2061,73 @@ model LiteLLM_TelemetryReport { @@index([window_end]) @@index([created_at]) } + +enum AccountingEventKind { + inbound_usage + outbound_usage + cost +} + +enum AccountingDeliveryState { + pending + applied + parked +} + +enum AccountingDeliveryTarget { + LiteLLM_VerificationToken + LiteLLM_UserTable + LiteLLM_TeamTable + LiteLLM_TeamMembership + LiteLLM_OrganizationTable + LiteLLM_OrganizationMembership + LiteLLM_EndUserTable + LiteLLM_AgentsTable + LiteLLM_TagTable + LiteLLM_ProjectTable + LiteLLM_ModelAccessGroupBudgetTable + LiteLLM_BudgetTable + LiteLLM_BudgetWindowSpend + LiteLLM_DailyUserSpend + LiteLLM_DailyTeamSpend + LiteLLM_DailyOrganizationSpend + LiteLLM_DailyEndUserSpend + LiteLLM_DailyAgentSpend + LiteLLM_DailyTagSpend + LiteLLM_DailyModelUsage + LiteLLM_DailyToolSpend + LiteLLM_AutoRouterDailySpend +} + +// Durable accounting ledger. One row per usage or cost event; payload holds the +// private event envelope versioned by schema_version. Usage rows retain request-scoped +// pricing context; cost rows include resolved rates and the finalized pricing snapshot. +model LiteLLM_AccountingEvent { + event_id String @id + kind AccountingEventKind + source_event_id String? + schema_version Int + payload_hash String + occurred_at DateTime + recorded_at DateTime @default(now()) + payload Json + + @@index([recorded_at]) + @@index([kind, source_event_id]) +} + +// One row per (event, target table, version): pending work for the projector and, +// once applied in the same transaction as the target write, the receipt that stops a retry. +model LiteLLM_AccountingDelivery { + event_id String + target AccountingDeliveryTarget + version Int + state AccountingDeliveryState @default(pending) + attempts Int @default(0) + next_attempt_at DateTime @default(now()) + applied_at DateTime? + last_error String? + + @@id([event_id, target, version]) + @@index([target, state, next_attempt_at]) +} diff --git a/skills/litellm-setup/SKILL.md b/skills/litellm-setup/SKILL.md new file mode 100644 index 00000000000..7c824751d6f --- /dev/null +++ b/skills/litellm-setup/SKILL.md @@ -0,0 +1,205 @@ +--- +name: litellm-setup +description: Set up a self-hosted LiteLLM proxy, the OpenAI-compatible LLM gateway, on this machine, from install through provider keys, config.yaml, a master key, optional Postgres and Admin UI, and a verified chat completion. Use when the user asks to set up LiteLLM, install the LiteLLM proxy, self-host an LLM gateway, or put one endpoint in front of OpenAI, Anthropic, Gemini, Azure, Bedrock, Ollama, or other providers. Not for using the LiteLLM Python SDK inside an app +--- + +# Set up a LiteLLM proxy + +Take the user from nothing to a running LiteLLM proxy that has answered one chat completion. Run every command yourself unless a step says the user does it + +Never print a secret. Provider keys, the master key, the salt key, the database URL, and virtual keys live only in `.env`, and `config.yaml` reads them through `os.environ/NAME`. Check that a value is there with `grep -c '^NAME=.' .env`, never with `cat .env` or `echo "$NAME"`. Load `.env` into a shell with this line, which exports each line as it is, so a `&` or `$` in a value stays literal, and first drops any database settings the shell already had and sets `LITELLM_MODE=PRODUCTION`, which stops `litellm` from also loading a `.env` it finds above its install folder, such as `~/.env`, so only this `.env` picks the database. The steps that need it start with it: + +```bash +unset DATABASE_URL DIRECT_URL DATABASE_HOST; export LITELLM_MODE=PRODUCTION; while IFS= read -r line || [ -n "$line" ]; do case "$line" in [A-Za-z_]*=*) export "$line" ;; esac; done < .env +``` + +## 1. Ask once + +Ask everything in one message, with defaults the user can accept by saying "go": + +- Providers and models. Default is OpenAI with its newest flagship model, using the `OPENAI_API_KEY` in your environment when `[ -n "$OPENAI_API_KEY" ] && echo set` says it is there. Run that check before asking, and say in the question that the default uses that variable, naming it but never its value, since a key in your environment can belong to another account +- Whether they want virtual keys, teams, spend tracking, or the Admin UI. Those need Postgres, plain routing does not. Default is no +- uv (default) or Docker. With Postgres, Docker runs the repo's quickstart, which brings its own database, and uv needs a Postgres server they already have, with a new, empty database for LiteLLM, so someone with their own Postgres server picks uv +- Folder and port. Defaults are `./litellm-proxy` (`./litellm-gateway` for the quickstart) and 4000 + +Then follow the steps for their answers. Plain routing skips steps 3 and 6. The quickstart is steps 2, 3, 8, and 9. uv with their own Postgres skips step 3 + +In the same message, say that someone who picks uv and would rather answer prompts can run `litellm --setup` in a separate terminal window once LiteLLM is installed, not through Claude Code. Never run it yourself, since it needs a real terminal and without one prints "Setup cancelled." and writes nothing. Give them the command as `LITELLM_MODE=PRODUCTION /litellm --setup`, with `` replaced by the folder `uv tool dir --bin` prints, to run from the folder they want the config in. It writes `litellm_config.yaml` there with the provider keys and master key in plain text and offers to start the proxy, which then listens on every network interface and stops with Ctrl+C in that window. When the user is done, run this in that folder. It makes the file readable only by them, keeps it out of git (the wizard does neither), and prints only the model names. Never print that file or run `grep`, `sed`, or `cat` on it: + +```bash +chmod 600 litellm_config.yaml; grep -qx 'litellm_config.yaml' .gitignore 2>/dev/null || echo 'litellm_config.yaml' >> .gitignore +"$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(*(m["model_name"] for m in yaml.safe_load(open("litellm_config.yaml"))["model_list"]), sep="\n")' +``` + +Then skip to step 8 and run it against their proxy, starting that shell command with this line instead of the loader line, which reads the master key without printing it: + +```bash +export LITELLM_MASTER_KEY="$("$(uv tool dir)/litellm/bin/python" -I -c 'import yaml; print(yaml.safe_load(open("litellm_config.yaml"))["general_settings"]["master_key"])')" +``` + +## 2. Check prerequisites + +```bash +uv --version; docker info --format '{{.ServerVersion}}'; docker compose version; curl --version | head -1; openssl version; command -v python3 +``` + +On native Windows, ask the user to run Claude Code inside WSL, since every command here is POSIX shell. Every path needs curl, openssl, and python3, so ask the user to install any that is missing. The uv path needs uv, which downloads a supported Python (LiteLLM supports 3.10 to 3.14) when the `--python` range in step 4 rules out the one on the machine. If uv is missing, offer to install it, and only after the user says yes run `curl -LsSf https://astral.sh/uv/install.sh | env UV_NO_MODIFY_PATH=1 sh`, which leaves their shell profile alone (https://docs.astral.sh/uv/getting-started/installation/). It puts uv in `~/.local/bin`, so start every later command that runs `uv` or `litellm` with `export PATH="$HOME/.local/bin:$PATH";`. The Docker path needs a running Docker daemon, and the quickstart also needs Docker Compose v2. If `docker info` says `permission denied`, the user's account cannot use Docker without sudo. Offer the uv path, or ask them to add themselves to the `docker` group the way https://docs.docker.com/engine/install/linux-postinstall/ shows and start a new login session before you check again + +Check the port is free with `curl -s -o /dev/null --max-time 2 http://127.0.0.1:4000; echo $?`, which prints `7` (could not connect) only when nothing listens there. Don't use `lsof` for this, since it misses ports that root or Docker hold. If it is taken, pick another port and use it everywhere below, because a proxy asked for a busy 4000 silently moves to a random port + +## 3. Postgres with Docker: run the quickstart + +The repo's quickstart starts the gateway, Postgres, and the Admin UI in Docker, listening on this machine only, and puts a `litellm-gateway` command in that folder that stops, starts, updates, or removes them. It asks nothing when it has no terminal. It writes the master key, salt key, and database password to its own `.env` and never prints them. Download it, then run it with the folder and port from step 1. It returns once the gateway is ready: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/scripts/quickstart.sh -o quickstart.sh +LITELLM_DIR="$PWD/litellm-gateway" LITELLM_PORT=4000 sh quickstart.sh --yes +cd litellm-gateway +``` + +Models live in the database here, so add each one through the API instead of a `config.yaml`, naming it the way step 5 says. Pipe the request body in, so the provider key never shows up in output or on a command line. Unless the user agreed to the key in your environment, as step 5 says, ask them to add the line `OPENAI_API_KEY=` to this `.env` themselves. Run the loader line first: + +```bash +python3 -c 'import json, os, sys; json.dump({"model_name": "gpt-6.1-sol", "litellm_params": {"model": "openai/gpt-6.1-sol", "api_key": os.environ["OPENAI_API_KEY"]}}, sys.stdout)' \ + | curl -sS -o /dev/null -w '%{http_code}\n' http://127.0.0.1:4000/model/new \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" -d @- +``` + +Run that in one shell command that starts with the loader line, since the master key is in `.env`. `200` means the model was added. Then go to step 8 + +## 4. Install + +uv: + +```bash +uv tool install --python '>=3.10,<3.15' 'litellm[proxy]' +export PATH="$(uv tool dir --bin):$PATH"; litellm --version +``` + +With their own Postgres, install `'litellm[proxy,extra_proxy]'` instead, with the same `--python` range, which adds the Prisma client the proxy uses for the database. Running `uv tool install` again with the new extras replaces the old install. Start every later command that runs `litellm` with `export PATH="$(uv tool dir --bin):$PATH";`, so this install runs even when the shell has no `litellm` or an older one from a venv or `pip install litellm`. `uv tool update-shell` puts it on the PATH of new shells by editing the user's shell profile, so run it only after the user says yes + +Docker: + +```bash +docker pull docker.litellm.ai/berriai/litellm:main-stable +``` + +## 5. Write `.env` and `config.yaml` + +Run every command from here on in the folder from step 1. Create `.env` readable only by the user, keep it out of git, and add a master key unless it already has one. The proxy refuses to start without one, and `sk-` is the format of every key it issues: + +```bash +mkdir -p litellm-proxy && cd litellm-proxy +touch .env && chmod 600 .env; grep -qx '.env' .gitignore 2>/dev/null || echo '.env' >> .gitignore +grep -q '^LITELLM_MASTER_KEY=.' .env || printf 'LITELLM_MASTER_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env +``` + +Copy a provider key that is already in your environment only after the user agreed to that variable, in their step 1 answer or when you ask now, and copy it without printing it, as in `printf 'OPENAI_API_KEY=%s\n' "$OPENAI_API_KEY" >> .env`. Otherwise ask the user to add the line `OPENAI_API_KEY=` to `.env` themselves and tell you when it is there. If they paste a key into the chat anyway, write it to `.env` and never repeat it. Keep one `NAME=value` per line with no quotes, so `docker --env-file` and the loader line read the same value + +Confirm each model's current LiteLLM name in the cost map instead of trusting memory, since model names go stale fast. This prints the chat models of the provider named last (`openai`, `anthropic`, `gemini`, `azure`, `bedrock`, `vertex_ai`, and so on) the way `config.yaml` takes them, with the provider in front. It matches by prefix, so `bedrock` also lists the newer Bedrock models and `vertex_ai` lists every Vertex family. For Ollama, use a model from `ollama list` instead: + +```bash +curl -fsSL https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json | python3 -c ' +import json, sys +for name, row in json.load(sys.stdin).items(): + provider = str(row.get("litellm_provider", "")) if isinstance(row, dict) else "" + if provider.startswith(sys.argv[1]) and row.get("mode") == "chat": + prefix = "bedrock" if provider == "bedrock_converse" else "vertex_ai" if provider.startswith("vertex_ai") else provider + print(name if name.startswith(prefix + "/") else prefix + "/" + name) +' openai +``` + +In `config.yaml`, `model` is a name from that list, as in `openai/gpt-6.1-sol` or `bedrock/us.anthropic.`, and `model_name` is what clients will send. If the folder already has a `config.yaml`, show the user its `model_name` lines and ask before replacing it. This example uses the newest GPT model in the cost map when this skill was written, so replace it with what the user picked: + +```yaml +model_list: + - model_name: gpt-6.1-sol + litellm_params: + model: openai/gpt-6.1-sol + api_key: os.environ/OPENAI_API_KEY + +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY +``` + +Add one entry per model. On Azure, ask the user for each deployment's name and the model it runs. Use `azure/` as `model`, since Azure routes by deployment and a name from the lookup works only when the deployment has that same name, and add `model_info:` with `base_model: azure/` to the entry, since LiteLLM prices each call by that name. Providers that need more than a key, like Azure or Bedrock, list their `litellm_params` on their page under https://docs.litellm.ai/docs/providers. Point every secret at `.env` through `os.environ/NAME` the same way + +## 6. Their own Postgres, on the uv path + +Ask the user to create a new, empty database for LiteLLM on their server (`CREATE DATABASE litellm;`) and add `DATABASE_URL=postgresql://...` for it to `.env` themselves. The proxy creates its tables there on first start and stops on a database that already holds other tables. If that URL goes through a transaction pooler such as PgBouncer, they also add `DIRECT_URL` for the same database without the pooler, which the migrations use. Then add a salt key. It encrypts the provider credentials the proxy stores in the database and falls back to the master key when unset, so set it once and never change it: + +```bash +grep -q '^LITELLM_SALT_KEY=.' .env || printf 'LITELLM_SALT_KEY=sk-%s\n' "$(openssl rand -hex 32)" >> .env +``` + +Generate the Prisma client once inside the tool's environment. A pip or uv install doesn't ship it generated (Docker images do): + +```bash +tool="$(uv tool dir)/litellm" +schema=$("$tool/bin/python" -I -c 'import importlib.util, os; print(os.path.join(importlib.util.find_spec("litellm").submodule_search_locations[0], "proxy", "schema.prisma"))') +PATH="$tool/bin:$PATH" prisma generate --schema "$schema" +``` + +## 7. Start the proxy + +Bind to `127.0.0.1` unless the user wants other machines to reach it + +uv, in one shell command that starts with the loader line, since the installed `litellm` does not read `.env` by itself: + +```bash +export PATH="$HOME/.local/bin:$PATH"; export PATH="$(uv tool dir --bin):$PATH" +nohup litellm --config config.yaml --host 127.0.0.1 --port 4000 > litellm.log 2>&1 & +echo $! > litellm.pid +``` + +Docker: + +```bash +docker run -d --name litellm --env-file .env -v "$PWD/config.yaml:/app/config.yaml:ro" -p 127.0.0.1:4000:4000 \ + docker.litellm.ai/berriai/litellm:main-stable --config /app/config.yaml --port 4000 +``` + +Wait on the health route, stopping early if the proxy exits. The first start with Postgres applies the database migrations and takes longer. Define `alive` for the path you took, `alive() { kill -0 "$(cat litellm.pid)" 2>/dev/null; }` for uv or `alive() { [ "$(docker inspect -f '{{.State.Running}}' litellm 2>/dev/null)" = true ]; }` for Docker, in the same shell command as the loop: + +```bash +for i in $(seq 1 180); do curl -fsS http://127.0.0.1:4000/health/liveliness >/dev/null 2>&1 && break; alive || break; sleep 1; done +curl -sS http://127.0.0.1:4000/health/readiness +``` + +Readiness reports `"db":"connected"` with Postgres and `"db":"Not connected"` without it. If the proxy never comes up, read the end of `litellm.log` (or `docker logs litellm`) and check the troubleshooting list + +## 8. Verify with a chat completion + +Use a `model_name` from `config.yaml` (or step 3), in one shell command that starts with the loader line: + +```bash +curl -sS http://127.0.0.1:4000/v1/chat/completions \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" -H "Content-Type: application/json" \ + -d '{"model": "gpt-6.1-sol", "messages": [{"role": "user", "content": "Say hello in five words."}]}' +``` + +Show the user the reply in `choices[0].message.content`. An error here means the setup is not done, so fix it before handing off + +## 9. Hand off + +Tell the user, without printing any secret: + +- The base URL is `http://127.0.0.1:4000`. Any OpenAI-compatible client works with it as the base URL and the master key or a virtual key as the API key +- The master key is `LITELLM_MASTER_KEY` in the folder's `.env`, or `general_settings.master_key` in `litellm_config.yaml` when they used `litellm --setup` +- To stop it, `kill "$(cat litellm.pid)"` (uv), `docker stop litellm` (Docker), `/litellm-gateway stop` (quickstart), or Ctrl+C in its window (`litellm --setup`). To start it again, `docker start litellm`, `/litellm-gateway start`, on uv the whole step 7 command as you ran it, loader line and PATH export included, wrapped in `( ... )` so their shell keeps its own settings, shown as a code block and never paraphrased, since a new shell has neither, or for `litellm --setup` `LITELLM_MODE=PRODUCTION /litellm --config litellm_config.yaml --host 127.0.0.1 --port ` in a terminal window in that folder, with `` and `` filled in, since running the wizard again writes a new config with a new master key. `update` and `uninstall` update or remove a quickstart install. Give the quickstart commands with the folder's full path, since a bare `litellm-gateway` is on PATH only when `~/.local/bin` is, and can belong to an earlier install +- With Postgres, the Admin UI is at `http://127.0.0.1:4000/ui` (user `admin`, password the master key), and `POST /key/generate` with the master key creates virtual keys with their own budgets and model lists. Without Postgres, both need step 3 or 6 first +- The Admin UI warns that the master-key login is a shared secret, which is fine while only this machine reaches the proxy. Before anyone else can reach it, invite an admin with their own password on the Internal Users page, then set `general_settings.disable_env_credential_login: true` in the proxy config +- To manage the running proxy from a coding agent (keys, teams, users, models, MCP servers), install `npx skills add BerriAI/litellm-skills -g` + +## Troubleshooting + +- `refused to start: no master key is set` means the loader line did not run in the shell command that started `litellm`, or `--env-file .env` is missing from `docker run` +- `a database URL is set but the prisma CLI is neither on PATH nor importable` means the install lacks `extra_proxy`, so rerun step 4 with it, or, on plain routing, that the command that started `litellm` did not start with the loader line +- `The Client hasn't been generated yet` or `Unable to find Prisma binaries` means the Prisma generate command in step 6 has not run +- `Incorrect API key provided` or another provider `AuthenticationError` means a provider key is wrong. Fix it in `.env` and restart, since the proxy reads the environment only at startup. On uv, stop it and rerun step 7. On Docker, `docker rm -f litellm` and rerun step 7, since `docker restart` keeps the environment the container was created with. On the quickstart the key is saved with the model, so delete the model and add it again: `GET /model/info` lists each model's `model_info.id`, and `POST /model/delete` with `{"id": ""}` deletes it +- `Cannot automatically baseline this database` means `DATABASE_URL` points at a database that already holds tables LiteLLM did not create. Point it at a new, empty database +- `LLM Provider NOT provided` means a `model` in `config.yaml` (or step 3) lacks its provider in front. Use the name as the step 5 lookup prints it +- `Invalid model name passed in model=...` means the request's `model` is not a `model_name` in `config.yaml` or step 3. `GET /v1/models` lists the ones the key can use +- `No connected db.` on a proxy without Postgres means the bearer token is not the master key +- `Conflict. The container name "/litellm" is already in use` means an earlier run left a container. `docker start litellm` reuses it, or with the user's OK `docker rm -f litellm` and rerun step 7 +- Inside Docker, `localhost` in `config.yaml` (an Ollama `api_base`, say) is the container itself. Use `host.docker.internal`, and on Linux also pass `--add-host=host.docker.internal:host-gateway` diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index e546549d25f..e37593adb4b 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -8,6 +8,8 @@ the very class of undeclared-dependency bug this guards against. import argparse import importlib.util +from importlib.machinery import EXTENSION_SUFFIXES +from pathlib import Path import re import sys import traceback @@ -89,6 +91,18 @@ def check_import() -> str: return f"imported {distribution} {sdk_version}" +def check_native_extension() -> str: + import litellm + from litellm.rust_bridge import _native + + native_path: Final = Path(_native.__file__).resolve() + package_path: Final = Path(litellm.__file__).resolve().parent + _require(native_path.parent == package_path / "rust_bridge", "native extension is outside the SDK package") + _require(any(str(native_path).endswith(suffix) for suffix in EXTENSION_SUFFIXES), "native module is not compiled") + _require(_native.trace_encode_error("ping") == b"\x12\x04ping", "native OTLP error encoding failed") + return f"loaded and exercised {native_path}" + + def check_completion() -> str: import litellm @@ -230,6 +244,7 @@ def check_bedrock_credential_resolution() -> str: CHECKS: tuple[tuple[str, Callable[[], str]], ...] = ( ("environment is base-only", check_environment_is_base_only), ("import litellm", check_import), + ("native extension", check_native_extension), ("optional MCP installation guidance", check_mcp_install_guidance), ("chat completion", check_completion), ("chat completion with a function tool", check_completion_with_function_tool), @@ -250,8 +265,11 @@ def _run(check: Callable[[], str]) -> tuple[bool, str]: def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--profile", choices=("legacy", "core", "dependencies"), default="legacy") + parser.add_argument("--profile", choices=("legacy", "core", "dependencies", "native"), default="legacy") profile = parser.parse_args().profile + if profile == "native": + print(check_native_extension()) + return 0 for module in ("boto3", "botocore", "tokenizers", "huggingface_hub"): present = importlib.util.find_spec(module) is not None _require(present == (profile != "core"), f"{profile}: unexpected presence of {module}: {present}") diff --git a/tests/benchmarks/test_sse_lines_benchmarks.py b/tests/benchmarks/test_sse_lines_benchmarks.py new file mode 100644 index 00000000000..1320fae49a6 --- /dev/null +++ b/tests/benchmarks/test_sse_lines_benchmarks.py @@ -0,0 +1,27 @@ +from collections.abc import Sequence +from typing import Final + +import pytest + +from litellm.llms.custom_httpx.sse_lines import iter_sse_lines + +EVENT: Final = 'data: {"choices":[{"delta":{"content":"hello world"}}]}' + + +def _chunks(line_break: str) -> Sequence[str]: + body: Final = "".join(f"{EVENT}{line_break}{line_break}" for _ in range(2000)) + return tuple(body[i : i + 512] for i in range(0, len(body), 512)) + + +LF_CHUNKS: Final = _chunks("\n") +CRLF_CHUNKS: Final = _chunks("\r\n") + + +@pytest.mark.benchmark +def test_iter_sse_lines_lf_stream() -> None: + assert sum(1 for _ in iter_sse_lines(LF_CHUNKS)) == 4000 + + +@pytest.mark.benchmark +def test_iter_sse_lines_crlf_stream() -> None: + assert sum(1 for _ in iter_sse_lines(CRLF_CHUNKS)) == 4000 diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index eee02a2f6b6..9fb3633d53c 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -175,6 +175,8 @@ ignored_function_names = [ "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py + "priority_header_update", # Tested through the priority header tests in test_common_request_processing.py + "_drops_params", # Read by split_request_priority and priority_header_update, tested through both "_responses_event_announces_background_run", # Tested through the background Responses stream tests in test_router.py "_aresponses_background_stream_can_release", # Tested through the background Responses stream tests in test_router.py "_raise_released_stream_error", # Tested through the background Responses stream content-policy tests in test_router.py @@ -185,13 +187,13 @@ ignored_function_names = [ def main(): router_file = [ "./litellm/router.py", - "./litellm/router_utils/batch_utils.py", + "./litellm/litellm_core_utils/batch_utils.py", "./litellm/router_utils/pattern_match_deployments.py", ] # router_file = [ # "../../litellm/router.py", # "../../litellm/router_utils/pattern_match_deployments.py", - # "../../litellm/router_utils/batch_utils.py", + # "../../litellm/litellm_core_utils/batch_utils.py", # ] ## LOCAL TESTING tests_dir = "./tests/" # Update this path if your tests directory is located elsewhere # tests_dir = "../../tests/" # LOCAL TESTING diff --git a/tests/code_coverage_tests/router_surface_allowlist.json b/tests/code_coverage_tests/router_surface_allowlist.json new file mode 100644 index 00000000000..428fd9b4011 --- /dev/null +++ b/tests/code_coverage_tests/router_surface_allowlist.json @@ -0,0 +1,211 @@ +{ + "_doc": "Router members and litellm.router_utils modules used outside the router. Shrink-only: see tests/code_coverage_tests/router_surface_check.py.", + "members_outside": [ + "_get_client", + "_zero_cost_cache", + "a_add_message", + "aadapter_completion", + "aanthropic_messages", + "abatch_completion", + "abatch_completion_fastest_response", + "acancel_batch", + "acancel_fine_tuning_job", + "acompletion", + "acreate_assistants", + "acreate_batch", + "acreate_file", + "acreate_fine_tuning_job", + "acreate_thread", + "adaptive_routers", + "adecisions", + "adelete_assistant", + "aembedding", + "afile_content", + "afile_delete", + "aget_assistants", + "aget_messages", + "aget_responses", + "aget_thread", + "alist_batches", + "alist_fine_tuning_jobs", + "allm_passthrough_route", + "amoderation", + "aocr", + "are_all_deployments_blocked", + "arefresh_model_info", + "aretrieve_batch", + "aretrieve_fine_tuning_job", + "arm_routing_read_prefetch", + "arun_thread", + "async_get_available_deployment", + "async_get_fully_unhealthy_model_names", + "async_get_healthy_deployments", + "async_pre_routing_hook", + "auto_routers", + "avector_store_create", + "avector_store_search", + "background_health_check_model_groups", + "bind_to_request_logging", + "cache", + "cache_responses", + "cached_deployment_model_info", + "cached_model_group_info", + "completion", + "complexity_routers", + "config_deployments", + "content_policy_fallbacks", + "context_window_fallbacks", + "cooldown_time", + "default_deployment", + "default_litellm_params", + "delete_deployment", + "deployment_latency_map", + "deployment_names", + "deployment_usable_by_team", + "deployments_for_request", + "deployments_named", + "discard", + "embedding", + "enable_health_check_routing", + "fallbacks", + "generate_model_id", + "get_all_deployments", + "get_available_deployment_for_pass_through", + "get_available_guardrail", + "get_candidate_model_ids_for_route", + "get_configured_display_name", + "get_configured_mode", + "get_configured_service_tiers", + "get_configured_token_limits", + "get_credential_deployment", + "get_deployment", + "get_deployment_by_model_group_name", + "get_deployment_credentials", + "get_deployment_credentials_with_provider", + "get_deployment_model_info", + "get_discovered_model_info", + "get_fully_blocked_model_names", + "get_model_access_groups", + "get_model_from_alias", + "get_model_group_info", + "get_model_group_usage", + "get_model_ids", + "get_model_info", + "get_model_list", + "get_model_list_from_model_alias", + "get_model_list_from_routing_groups", + "get_model_list_of_routed_group", + "get_model_listing_info", + "get_model_names", + "get_remaining_model_group_usage", + "get_routable_deployments", + "get_routable_upstream_model", + "get_router_deployment_budget_limiter", + "get_routing_group", + "get_settings", + "get_valid_args", + "guardrail_list", + "has_model_id", + "health_check_ignore_transient_errors", + "health_state_cache", + "is_deployment_blocked", + "is_model_access_group_for_wildcard_route", + "is_recognized_model", + "kubernetes_pod_discovery", + "map_team_model", + "model_group_alias", + "model_list", + "model_name_to_deployment_indices", + "model_names", + "pattern_router", + "priority_header_update", + "prompt_caching_pin_cache", + "provider_budget_config", + "quality_routers", + "resolve_model_name_from_model_id", + "resolve_served_model_name", + "resolved_litellm_models", + "routable_model_group", + "router_general_settings", + "routing_groups", + "routing_plugins", + "schedule_acompletion", + "search_tools", + "serve_through_request_logging", + "should_include_deployment", + "split_request_priority", + "team_model_to_deployment_indices", + "team_pattern_routers", + "team_ptu_ceiling", + "team_public_model_names", + "team_public_model_names_of", + "update_redis_cache", + "update_settings", + "upsert_deployment", + "wildcard_route_deployments", + "wildcard_routed_models" + ], + "members_router_family": [ + "_count_pre_call_check_tokens", + "_pre_call_checks", + "_update_kwargs_before_fallbacks", + "acompletion", + "adecisions", + "aembedding", + "allowed_fails", + "allowed_fails_policy", + "aresponses", + "async_function_with_fallbacks", + "async_get_healthy_deployments", + "cache", + "cached_deployment_model_info", + "cooldown_cache", + "cooldown_time", + "default_max_parallel_requests", + "deployments_named", + "disable_cooldowns", + "embedding", + "enable_pre_call_checks", + "enable_tag_filtering", + "explicit_allowed_fails", + "fallback_access_check", + "fallback_budget_check", + "fallbacks", + "get_all_deployments", + "get_allowed_fails_from_policy", + "get_deployment", + "get_model_group", + "get_model_ids", + "get_model_info", + "get_model_list", + "get_router_model_info", + "log_retry", + "model_list", + "provider_default_deployment_ids", + "routing_group_has_alternatives", + "routing_strategy", + "search_tools", + "set_custom_routing_strategy", + "slack_alerting_logger", + "tag_filtering_match_any", + "tag_routing_prefix", + "team_model_has_alternatives" + ], + "router_utils_imports_outside": [ + "litellm.router_utils.access_windows", + "litellm.router_utils.add_retry_fallback_headers", + "litellm.router_utils.clientside_credential_handler", + "litellm.router_utils.common_utils", + "litellm.router_utils.cooldown_handlers", + "litellm.router_utils.deployment_pricing", + "litellm.router_utils.fallback_event_handlers", + "litellm.router_utils.get_retry_from_policy", + "litellm.router_utils.health_state_cache", + "litellm.router_utils.pattern_match_deployments", + "litellm.router_utils.ptu_shares", + "litellm.router_utils.reasoning_effort_capability", + "litellm.router_utils.router_callbacks.track_deployment_metrics", + "litellm.router_utils.routing_groups", + "litellm.router_utils.search_api_router" + ] +} diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py new file mode 100644 index 00000000000..0b35bac5240 --- /dev/null +++ b/tests/code_coverage_tests/router_surface_check.py @@ -0,0 +1,235 @@ +""" +Fence for the Router's public surface. + +Code outside the router may use only the Router members listed in +router_surface_allowlist.json, and may import only the litellm.router_utils +modules listed there. The lists hold what is used today and may only shrink: + +- a member or module used outside the router that is not listed fails the + check (use an existing entry point instead, or discuss widening the surface); +- a listed member or module that nothing outside the router uses any more + also fails, so the entry is removed in the same change. + +"Outside" means every non-test file under litellm/ and enterprise/ except +litellm/router.py and the router's own modules. Uses from the router family +(litellm/router_utils, litellm/router_strategy) are listed separately, because +those modules become private parts of the router core in a later step. + +Member uses are found the same way the boundary survey found them: attribute +access on a receiver whose name looks like a router (llm_router, +litellm_router_instance, router, ...) or on the Router class, plus +getattr/setattr/hasattr with a constant name on such a receiver. Members +include the container file endpoints Router.__init__ registers by name. + +Run: python tests/code_coverage_tests/router_surface_check.py [--write] +--write regenerates the allowlist from the current tree (for shrinking it). +""" + +import ast +import json +import os +import re +import sys +from typing import Dict, List, Set + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +ALLOWLIST = os.path.join(os.path.dirname(__file__), "router_surface_allowlist.json") + +ROUTER_FILES = { + "litellm/router.py", +} +ROUTER_CLASS_FILES = ["litellm/router.py"] +ROUTER_CLASS_NAMES = {"Router"} + +SCAN_ROOTS = ["litellm", "enterprise"] +SKIP_DIRS = { + ".git", + "node_modules", + ".venv", + "__pycache__", + "litellm-rust", + "ui", + "docs", + "cookbook", + "tests", + "test", +} + +RECEIVER = re.compile( + r"(^|[._])(llm_router|router|litellm_router|litellm_router_instance|router_instance" + r"|_router|proxy_router|llm_router_instance|router_obj|mock_router|Router)$" +) +# Router.__init__ registers these with setattr(self, name, ...), so they never appear as self.. +CONTAINER_FILE_ENDPOINTS = "litellm/containers/endpoints.json" +DYNAMIC_FUNCS = {"getattr", "setattr", "hasattr"} + + +def _area(rel: str) -> str: + if rel.startswith(("litellm/router_utils/", "litellm/router_strategy/")): + return "router_family" + return "outside" + + +def _is_test(rel: str) -> bool: + name = os.path.basename(rel) + return "/tests/" in rel or name.startswith("test_") or name.endswith("_test.py") + + +def router_members() -> Set[str]: + members: Set[str] = set() + for rel in ROUTER_CLASS_FILES: + with open(os.path.join(ROOT, rel), encoding="utf-8") as f: + tree = ast.parse(f.read()) + for cls in tree.body: + if not isinstance(cls, ast.ClassDef) or cls.name not in ROUTER_CLASS_NAMES: + continue + for node in cls.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + members.add(node.name) + elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): + members.add(node.target.id) + elif isinstance(node, ast.Assign): + members.update(t.id for t in node.targets if isinstance(t, ast.Name)) + for node in ast.walk(cls): + if ( + isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "self" + and isinstance(node.ctx, ast.Store) + ): + members.add(node.attr) + with open(os.path.join(ROOT, CONTAINER_FILE_ENDPOINTS), encoding="utf-8") as f: + for endpoint in json.load(f)["endpoints"]: + members.update((endpoint["name"], endpoint["async_name"])) + return {m for m in members if not (m.startswith("__") and m.endswith("__"))} + + +def _python_files(): + for top in SCAN_ROOTS: + base = os.path.join(ROOT, top) + for dirpath, dirnames, filenames in os.walk(base): + dirnames[:] = [d for d in dirnames if d not in SKIP_DIRS] + for fn in filenames: + if fn.endswith(".py"): + path = os.path.join(dirpath, fn) + yield os.path.relpath(path, ROOT).replace(os.sep, "/"), path + + +def scan() -> Dict[str, Dict[str, Set[str]]]: + """Return {"members": {area: {member: files}}, "router_utils_imports": {module: files}}.""" + members = router_members() + uses: Dict[str, Dict[str, Set[str]]] = {"outside": {}, "router_family": {}} + imports: Dict[str, Set[str]] = {} + for rel, path in _python_files(): + if rel in ROUTER_FILES or _is_test(rel): + continue + with open(path, encoding="utf-8") as f: + src = f.read() + has_router = "router" in src.lower() + if not has_router: + continue + try: + tree = ast.parse(src) + except SyntaxError: + continue + area = _area(rel) + for name in member_uses(tree, members): + uses[area].setdefault(name, set()).add(rel) + if area == "outside": + for node in ast.walk(tree): + for mod in _router_utils_modules(node): + imports.setdefault(mod, set()).add(rel) + return {"members": uses, "router_utils_imports": imports} + + +def member_uses(tree: ast.AST, members: Set[str]) -> Set[str]: + """The Router members a module uses, on an instance or on the Router class itself.""" + used: Set[str] = set() + for node in ast.walk(tree): + if isinstance(node, ast.Attribute) and node.attr in members: + if RECEIVER.search(ast.unparse(node.value)): + used.add(node.attr) + elif ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id in DYNAMIC_FUNCS + and len(node.args) >= 2 + and isinstance(node.args[1], ast.Constant) + and node.args[1].value in members + and RECEIVER.search(ast.unparse(node.args[0])) + ): + used.add(node.args[1].value) + return used + + +def _router_utils_modules(node: ast.AST) -> List[str]: + if isinstance(node, ast.ImportFrom) and node.module: + if node.module == "litellm.router_utils": + return ["litellm.router_utils." + a.name for a in node.names] + if node.module.startswith("litellm.router_utils."): + return [node.module] + if isinstance(node, ast.Import): + return [a.name for a in node.names if a.name.startswith("litellm.router_utils.")] + return [] + + +def build_allowlist(found) -> dict: + return { + "_doc": ( + "Router members and litellm.router_utils modules used outside the router. " + "Shrink-only: see tests/code_coverage_tests/router_surface_check.py." + ), + "members_outside": sorted(found["members"]["outside"]), + "members_router_family": sorted(found["members"]["router_family"]), + "router_utils_imports_outside": sorted(found["router_utils_imports"]), + } + + +def compare(allow: dict, found) -> List[str]: + errors: List[str] = [] + sections = [ + ("members_outside", found["members"]["outside"], "Router member"), + ("members_router_family", found["members"]["router_family"], "Router member (router family)"), + ("router_utils_imports_outside", found["router_utils_imports"], "litellm.router_utils import"), + ] + for key, used, what in sections: + listed = set(allow.get(key, [])) + for name in sorted(set(used) - listed): + files = ", ".join(sorted(used[name])) + errors.append( + f"new {what} outside the router: {name!r} in {files}. The router surface may only " + f"shrink: use an entry point instead." + ) + for name in sorted(listed - set(used)): + errors.append( + f"{what} {name!r} is listed in {os.path.basename(ALLOWLIST)} [{key}] but nothing uses it " + f"any more: remove it from the list." + ) + return errors + + +def main(argv: List[str]) -> int: + found = scan() + if "--write" in argv: + with open(ALLOWLIST, "w", encoding="utf-8") as f: + json.dump(build_allowlist(found), f, indent=2) + f.write("\n") + print(f"wrote {ALLOWLIST}") + return 0 + with open(ALLOWLIST, encoding="utf-8") as f: + allow = json.load(f) + errors = compare(allow, found) + for e in errors: + print(e) + if errors: + return 1 + print( + f"router surface ok: {len(allow['members_outside'])} members outside, " + f"{len(allow['members_router_family'])} from the router family, " + f"{len(allow['router_utils_imports_outside'])} router_utils modules imported outside" + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/tests/code_coverage_tests/unbounded_in_baseline.txt b/tests/code_coverage_tests/unbounded_in_baseline.txt index 1da964681a9..d8c6e2241f4 100644 --- a/tests/code_coverage_tests/unbounded_in_baseline.txt +++ b/tests/code_coverage_tests/unbounded_in_baseline.txt @@ -129,4 +129,4 @@ litellm/proxy/utils.py prefetch_config_params prisma param_name.in `param_names` litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma [scope.entity_id_field].in `list(scope.entity_ids)` 0 litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma api_key.in `list(scope.api_keys)` 0 litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma not.in `list(scope.exclude_entity_ids)` 0 -litellm/router_utils/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0 +litellm/router_strategy/auto_router/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0 diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index b2c5b1c74b6..005e581c643 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -254,8 +254,13 @@ _ROUTES: Final[dict[str, _Route]] = { "acreate_agent": ( "/v1beta/agents", "/v1beta/agents", - ("name",), - {"name": "agent", "base_agent": "waverunner", "instructions": "You are a helpful assistant."}, + ("id",), + { + "id": "agent", + "base_agent": "antigravity-preview-09-2026", + "system_instruction": "You are a helpful assistant.", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + }, "gemini/gemini-2.5-flash", ), "acreate_interaction": ( @@ -1250,11 +1255,7 @@ def test_valid_agent_creation_reaches_upstream(gateway: Gateway) -> None: _assert_scripted_response("acreate_agent", caller) matches: Final = _observed(_Observations(gateway.upstream_url), identity) outbound: Final = object_value(matches[0]["body"]) - assert outbound == { - "name": "agent", - "base_agent": "waverunner", - "instructions": "You are a helpful assistant.", - }, matches + assert outbound == body, matches def test_valid_speech_returns_binary_audio_and_reaches_upstream(gateway: Gateway) -> None: diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 1de2090b487..f558c355c84 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -1,3 +1,4 @@ +import hashlib import os import queue import secrets @@ -973,6 +974,7 @@ def test_the_migration_job_refuses_before_the_first_migration_when_spend_log_key _insert_spend_log_with_key(conn, "legacy-1", _legacy_bearer_token()) _insert_spend_log_with_key(conn, "legacy-2", _legacy_bearer_token()) _insert_spend_log_with_key(conn, "hashed-1", "a" * 64) + _insert_spend_log_with_key(conn, "repetitive-1", "z" * 3000) ledger_before: Final = _ledger(scratch_database) with pytest.raises(RuntimeError) as refused: @@ -1109,6 +1111,68 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( assert finding.detail == f"with an index entry over 2704 bytes, largest 2712 bytes ({described_as})" +@requires_db +def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random_text_of_that_length_is( + scratch_database: str, +) -> None: + """Postgres compresses a long value before it checks the entry size, so 3000 repeated + bytes index fine while 3000 random hex bytes do not. The check must draw the same line.""" + create_index: Final = sql.SQL('CREATE INDEX "compress_probe_idx" ON "compress_probe" ("k")') + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute('CREATE TABLE "compress_probe" ("k" TEXT)') + conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s), (%s)', ("a" * 3000, secrets.token_hex(1500))) + findings: Final = migration_data_checks.find_blockers( + conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe") + ) + with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"): + conn.execute(create_index) + conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,)) + conn.execute(create_index) + stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone() + + assert stored is not None and stored[0] == 3000 and stored[1] < 2704, stored + (finding,) = findings + assert isinstance(finding, migration_data_checks.BlockedChange) + assert finding.rows == 1, "only the random key is too long for the index once Postgres compresses the other" + assert finding.detail == "with an index entry over 2704 bytes, largest 3016 bytes (k up to 3000 bytes)" + + +def _postgres_rejects_an_index_on(conn: "psycopg.Connection[tuple[object, ...]]", key: str) -> bool: + conn.execute('DROP TABLE IF EXISTS "one_key"') + conn.execute('CREATE TABLE "one_key" ("k" TEXT)') + conn.execute('INSERT INTO "one_key" ("k") VALUES (%s)', (key,)) + try: + conn.execute('CREATE INDEX "one_key_idx" ON "one_key" ("k")') + except psycopg.errors.ProgramLimitExceeded: + return True + return False + + +@requires_db +def test_the_check_flags_exactly_the_compressed_keys_postgres_rejects_at_the_entry_limit(scratch_database: str) -> None: + """Keys that compress to a stored size just under and just over the entry limit go through the + check and through Postgres one by one. The check must flag a key exactly when Postgres rejects it.""" + noise: Final = "".join(hashlib.sha256(str(i).encode()).hexdigest() for i in range(40)) + keys: Final = tuple(noise[:n] + "a" * 1500 for n in range(2400, 2500)) + create_index: Final = sql.SQL('CREATE INDEX "stored_probe_idx" ON "stored_probe" ("k")') + with psycopg.connect(scratch_database, autocommit=True) as conn: + conn.execute('CREATE TABLE "stored_probe" ("k" TEXT)') + conn.execute('INSERT INTO "stored_probe" ("k") SELECT unnest(%s::text[])', (list(keys),)) + compressed: Final = conn.execute( + 'SELECT count(*) FROM "stored_probe" WHERE pg_column_size("k") < octet_length("k") AND octet_length("k") > 2704' + ).fetchone() + findings: Final = migration_data_checks.find_blockers( + conn, "public", migration_data_checks.checks_for_statement(create_index.as_string(conn), "probe") + ) + rejected: Final = tuple(key for key in keys if _postgres_rejects_an_index_on(conn, key)) + + assert compressed == (len(keys),), "every probe key must be longer than the limit yet stored compressed" + assert 0 < len(rejected) < len(keys), f"the probe must straddle the limit, Postgres rejected {len(rejected)}" + (finding,) = findings + assert isinstance(finding, migration_data_checks.BlockedChange) + assert finding.rows == len(rejected), f"the check must flag the {len(rejected)} keys Postgres rejects and no other" + + @requires_db def test_a_scan_that_cannot_finish_in_time_is_reported_as_unverified_rather_than_passed(scratch_database: str) -> None: with psycopg.connect(scratch_database, autocommit=True) as conn: diff --git a/tests/integration/database/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py new file mode 100644 index 00000000000..0c2ceb62eaf --- /dev/null +++ b/tests/integration/database/test_accounting_ledger.py @@ -0,0 +1,534 @@ +from __future__ import annotations + +import asyncio +import json +import os +import subprocess +import sys +from collections.abc import AsyncGenerator, AsyncIterator, Iterator, Mapping, Sequence +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from decimal import Decimal +from itertools import chain +from pathlib import Path +from typing import TYPE_CHECKING, Final, TypeAlias + +import psycopg +import pytest +import pytest_asyncio +from psycopg import sql +from pydantic import JsonValue + +from litellm.accounting import ( + CostEvent, + DeliveryState, + DeliveryTarget, + EventKind, + InboundUsageEvent, + Metric, + MetricName, + OutboundUsageEvent, + Outcome, + PricingSnapshot, +) +from litellm.proxy.db import accounting_ledger +from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger +from litellm.types.utils import CallTypes +from tests.integration._support.database import scratch_database +from tests.integration._support.prisma_clients import connected_proxy_prisma + +if TYPE_CHECKING: + from prisma import Prisma + + from litellm.proxy.utils import PrismaClient + +_Writer: TypeAlias = tuple["PrismaClient", AccountingLedger] + +_SCHEMA: Final = Path(__file__).resolve().parents[3] / "litellm-proxy-extras/litellm_proxy_extras/schema.prisma" +_EVENT_INSERT: Final = ( + 'INSERT INTO "LiteLLM_AccountingEvent" ' + "(event_id, kind, schema_version, payload_hash, occurred_at, payload) " + "VALUES ('event', %s, %s, 'hash', '2026-01-01', '{}') RETURNING kind::text, schema_version" +) +_DELIVERY_INSERT: Final = ( + 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version, state, attempts) ' + "VALUES ('event', %s, %s, %s, %s) RETURNING target::text, version, state::text, attempts" +) + + +@pytest.fixture(scope="module") +def ledger_database() -> Iterator[str]: + with scratch_database("accounting") as database_url: + result: Final = subprocess.run( + [sys.executable, "-I", "-m", "prisma", "migrate", "deploy", "--schema", str(_SCHEMA)], + env={**os.environ, "DATABASE_URL": database_url}, + capture_output=True, + text=True, + timeout=300, + ) + assert result.returncode == 0, result.stdout + result.stderr + yield database_url + + +@pytest.fixture +def ledger(ledger_database: str) -> Iterator[psycopg.Connection[tuple[object, ...]]]: + with psycopg.connect(ledger_database) as connection: + try: + yield connection + finally: + connection.rollback() + + +@pytest.mark.parametrize("kind", tuple(EventKind)) +def test_ledger_accepts_every_public_event_kind( + ledger: psycopg.Connection[tuple[object, ...]], kind: EventKind +) -> None: + row: Final = ledger.execute(_EVENT_INSERT, (kind.value, 1)).fetchone() + assert row == (kind.value, 1) + + +@pytest.mark.parametrize("target", tuple(DeliveryTarget)) +@pytest.mark.parametrize("state", tuple(DeliveryState)) +def test_delivery_accepts_every_public_target_and_state( + ledger: psycopg.Connection[tuple[object, ...]], target: DeliveryTarget, state: DeliveryState +) -> None: + row: Final = ledger.execute(_DELIVERY_INSERT, (target.value, 1, state.value, 0)).fetchone() + assert row == (target.value, 1, state.value, 0) + + +@pytest.mark.parametrize( + ("enum_name", "values"), + [ + ("AccountingEventKind", tuple(value.value for value in EventKind)), + ("AccountingDeliveryState", tuple(value.value for value in DeliveryState)), + ("AccountingDeliveryTarget", tuple(value.value for value in DeliveryTarget)), + ], +) +def test_migrated_enums_match_the_public_contract( + ledger: psycopg.Connection[tuple[object, ...]], enum_name: str, values: tuple[str, ...] +) -> None: + rows: Final = ledger.execute( + "SELECT e.enumlabel FROM pg_enum e JOIN pg_type t ON t.oid = e.enumtypid WHERE t.typname = %s", + (enum_name,), + ).fetchall() + assert frozenset(row[0] for row in rows) == frozenset(values) + + +@pytest.mark.parametrize("version", [0, -1]) +def test_event_rejects_nonpositive_schema_versions( + ledger: psycopg.Connection[tuple[object, ...]], version: int +) -> None: + with pytest.raises(psycopg.errors.CheckViolation, match="AccountingEvent_schema_version_check"): + ledger.execute(_EVENT_INSERT, (EventKind.OUTBOUND_USAGE.value, version)) + + +@pytest.mark.parametrize( + ("version", "attempts", "constraint"), + [ + (0, 0, "AccountingDelivery_version_check"), + (-1, 0, "AccountingDelivery_version_check"), + (1, -1, "AccountingDelivery_attempts_check"), + ], +) +def test_delivery_rejects_invalid_versions_and_attempts( + ledger: psycopg.Connection[tuple[object, ...]], version: int, attempts: int, constraint: str +) -> None: + with pytest.raises(psycopg.errors.CheckViolation, match=constraint): + ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, version, "pending", attempts)) + + +@pytest.mark.parametrize("column", ["state", "target"]) +def test_delivery_rejects_unknown_enum_values(ledger: psycopg.Connection[tuple[object, ...]], column: str) -> None: + ledger.execute(_DELIVERY_INSERT, (DeliveryTarget.KEY_SPEND.value, 1, "pending", 0)) + with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): + ledger.execute( + sql.SQL('UPDATE "LiteLLM_AccountingDelivery" SET {} = %s').format(sql.Identifier(column)), ("invalid",) + ) + + +def test_event_rejects_unknown_kind(ledger: psycopg.Connection[tuple[object, ...]]) -> None: + with pytest.raises(psycopg.errors.InvalidTextRepresentation, match="invalid input value for enum"): + ledger.execute(_EVENT_INSERT, ("invalid", 1)) + + +def test_new_delivery_defaults_to_pending_with_no_attempts(ledger: psycopg.Connection[tuple[object, ...]]) -> None: + row: Final = ledger.execute( + 'INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version) ' + "VALUES ('event', %s, 1) RETURNING state::text, attempts", + (DeliveryTarget.KEY_SPEND.value,), + ).fetchone() + assert row == ("pending", 0) + + +_AT: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +_CUTOFF: Final = datetime(2026, 2, 1, tzinfo=timezone.utc) +_TARGETS: Final = (DeliveryTarget.USER_SPEND, DeliveryTarget.TEAM_SPEND) + + +def _usage(event_id: str = "usage", *, inbound: bool = False) -> InboundUsageEvent | OutboundUsageEvent: + fields: Final = dict( + event_id=event_id, + request_id="request", + call_type=CallTypes.completion, + started_at=_AT, + ended_at=_AT, + outcome=Outcome.SUCCEEDED, + dimensions=(), + ) + if inbound: + return InboundUsageEvent(**fields, metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),)) + return OutboundUsageEvent(**fields, attempt_id=event_id, retry_of_attempt_id=None, metrics=None) + + +def _cost(source: str = "usage", event_id: str = "cost", amount: str = "0.125") -> CostEvent: + return CostEvent( + event_id=event_id, + usage_event_id=source, + occurred_at=_AT, + dimensions=(), + charges=(Metric(name=MetricName.SPEND, value=Decimal(amount), unit="USD"),), + snapshot=PricingSnapshot( + calculator_version="fixture", + model="fixture", + provider=None, + pricing_source="cost_map", + priced_model_key="fixture", + total_cost=Decimal(amount), + ), + ) + + +def _envelope(event: AccountingEvent, *, context: Mapping[str, JsonValue] | None = None) -> str: + return json.dumps({"event": event.model_dump(mode="json"), "context": context or {}}) + + +async def _insert( + store: AccountingLedger, + event: AccountingEvent, + targets: tuple[DeliveryTarget, ...] = _TARGETS, + *, + context: Mapping[str, JsonValue] | None = None, +) -> str: + return await store.insert_event(event, targets, envelope=_envelope(event, context=context)) + + +def _rows(database_url: str, statement: str, parameters: tuple[object, ...] = ()) -> Sequence[tuple[object, ...]]: + with psycopg.connect(database_url) as connection: + cursor: Final = connection.execute(statement, parameters) + return cursor.fetchall() if cursor.description is not None else [] + + +def _complete_and_age(database_url: str) -> None: + with psycopg.connect(database_url) as connection: + connection.execute('UPDATE "LiteLLM_AccountingEvent" SET recorded_at = %s', (_AT,)) + connection.execute("UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'applied', applied_at = %s", (_AT,)) + + +@pytest_asyncio.fixture(loop_scope="function") +async def writer(ledger_database: str, monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_Writer]: + monkeypatch.setenv("DATABASE_URL", ledger_database) + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + _rows(ledger_database, 'TRUNCATE "LiteLLM_AccountingDelivery", "LiteLLM_AccountingEvent"') + async with connected_proxy_prisma(ledger_database) as client: + yield client, AccountingLedger.for_client(client) + + +@pytest.mark.asyncio +async def test_event_and_complete_delivery_set_commit_together(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage() + assert await _insert(store, event, context={"requested_region": "eu", "quantity": "1.25"}) == "inserted" + rows: Final = _rows(ledger_database, 'SELECT kind::text, source_event_id, payload FROM "LiteLLM_AccountingEvent"') + assert rows == [ + ("outbound_usage", None, json.loads(_envelope(event, context={"requested_region": "eu", "quantity": "1.25"}))) + ] + deliveries: Final = _rows( + ledger_database, 'SELECT target::text, state::text, attempts FROM "LiteLLM_AccountingDelivery"' + ) + assert set(deliveries) == {(target.value, "pending", 0) for target in _TARGETS} + + +@pytest.mark.asyncio +async def test_concurrent_duplicates_preserve_completed_deliveries(writer: _Writer, ledger_database: str) -> None: + _, store = writer + event: Final = _usage(inbound=True) + results: Final = await asyncio.gather(*(_insert(store, event) for _ in range(8))) + assert results.count("inserted") == 1 and results.count("duplicate") == 7 + _complete_and_age(ledger_database) + before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') + reordered: Final = json.dumps(json.loads(_envelope(event)), sort_keys=True, indent=2) + assert await store.insert_event(event, tuple(reversed(_TARGETS)) + _TARGETS, envelope=reordered) == "duplicate" + assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY target') == before + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["payload", "context", "targets", "schema_version", "version"]) +async def test_conflicting_acceptance_cannot_overwrite_or_add_work( + writer: _Writer, ledger_database: str, change: str +) -> None: + _, store = writer + event: Final = _usage() + await _insert(store, event) + altered: Final = event.model_copy(update={"outcome": Outcome.FAILED}) if change == "payload" else event + targets: Final = (DeliveryTarget.KEY_SPEND,) if change == "targets" else _TARGETS + context: Final = {"override": "0.5"} if change == "context" else {} + options: Final = {change: 2} if change in ("schema_version", "version") else {} + with pytest.raises(AccountingConflictError, match="conflicts"): + await store.insert_event(altered, targets, envelope=_envelope(altered, context=context), **options) + assert _rows(ledger_database, 'SELECT payload FROM "LiteLLM_AccountingEvent"') == [(json.loads(_envelope(event)),)] + assert set(_rows(ledger_database, 'SELECT target::text FROM "LiteLLM_AccountingDelivery"')) == { + (t.value,) for t in _TARGETS + } + + +@pytest.mark.asyncio +async def test_competing_cost_identities_accept_only_one_authoritative_charge( + writer: _Writer, ledger_database: str +) -> None: + _, store = writer + await _insert(store, _usage()) + results: Final = await asyncio.gather( + _insert(store, _cost()), _insert(store, _cost(event_id="other-cost")), return_exceptions=True + ) + assert sum(result == "inserted" for result in results) == 1 + assert sum(isinstance(result, AccountingConflictError) for result in results) == 1 + accepted: Final = _rows( + ledger_database, "SELECT event_id, payload FROM \"LiteLLM_AccountingEvent\" WHERE kind = 'cost'" + ) + assert len(accepted) == 1 + event: Final = CostEvent.model_validate_json(json.dumps(accepted[0][1]["event"])) + assert await _insert(store, event) == "duplicate" + assert event.snapshot.total_cost == Decimal("0.125") + with pytest.raises(AccountingConflictError): + await _insert(store, _cost(event_id=event.event_id, amount="0.5")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("source", ["missing", "inbound", "cost"]) +async def test_cost_requires_an_existing_outbound_source(writer: _Writer, ledger_database: str, source: str) -> None: + _, store = writer + if source == "inbound": + await _insert(store, _usage("source", inbound=True)) + if source == "cost": + await _insert(store, _usage()) + await _insert(store, _cost(event_id="source")) + with pytest.raises(AccountingConflictError, match="outbound usage"): + await _insert(store, _cost(source="source", event_id="invalid")) + assert _rows(ledger_database, "SELECT count(*) FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'invalid'") == [ + (0,) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "cancellation"]) +async def test_interruption_rolls_back_or_retry_recognizes_commit( + writer: _Writer, ledger_database: str, failure: str +) -> None: + client, store = writer + + @asynccontextmanager + async def interrupted() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + if failure == "before_commit": + raise RuntimeError("injected pre-commit failure") + if failure == "cancellation": + raise asyncio.CancelledError + raise RuntimeError("injected lost acknowledgement") + + error: Final = asyncio.CancelledError if failure == "cancellation" else RuntimeError + with pytest.raises(error): + await _insert(AccountingLedger(interrupted), _usage()) + expected: Final = 1 if failure == "after_commit" else 0 + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(expected,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(expected * len(_TARGETS),)] + assert await _insert(store, _usage()) == ("duplicate" if expected else "inserted") + + +@pytest.mark.asyncio +async def test_delivery_insert_failure_rolls_back_event(writer: _Writer, ledger_database: str) -> None: + from prisma.errors import RawQueryError + + _, store = writer + with psycopg.connect(ledger_database) as connection: + connection.execute( + 'ALTER TABLE "LiteLLM_AccountingDelivery" ADD CONSTRAINT reject_delivery CHECK (version < 1)' + ) + try: + with pytest.raises(RawQueryError, match="reject_delivery"): + await _insert(store, _usage()) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] + finally: + with psycopg.connect(ledger_database) as connection: + connection.execute('ALTER TABLE "LiteLLM_AccountingDelivery" DROP CONSTRAINT reject_delivery') + + +@pytest.mark.asyncio +@pytest.mark.parametrize("blocked", ["pending", "parked", "unpriced", "recent_cost", "recovery"]) +async def test_retention_preserves_incomplete_or_recovery_needed_families( + writer: _Writer, ledger_database: str, blocked: str +) -> None: + _, store = writer + await _insert(store, _usage()) + if blocked != "unpriced": + await _insert(store, _cost()) + _complete_and_age(ledger_database) + with psycopg.connect(ledger_database) as connection: + if blocked in ("pending", "parked"): + connection.execute( + "UPDATE \"LiteLLM_AccountingDelivery\" SET state = %s WHERE event_id = 'cost'", (blocked,) + ) + if blocked == "recent_cost": + connection.execute( + "UPDATE \"LiteLLM_AccountingEvent\" SET recorded_at = %s WHERE event_id = 'cost'", (_CUTOFF,) + ) + before: Final = _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') + assert ( + await store.purge_completed( + redelivery_before=_CUTOFF, recovery_before=_AT if blocked == "recovery" else _CUTOFF + ) + == 0 + ) + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == before + + +@pytest.mark.asyncio +async def test_retention_removes_completed_family_and_receipts_together(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage()) + await _insert(store, _cost()) + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(0,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] + + +@pytest.mark.asyncio +async def test_retention_respects_batch_bound_and_skips_locked_families(writer: _Writer, ledger_database: str) -> None: + client, store = writer + for name in ("a", "b", "c"): + await _insert(store, _usage(name, inbound=True)) + _complete_and_age(ledger_database) + async with client.writer_db.tx() as tx: + await tx.query_raw("SELECT event_id FROM \"LiteLLM_AccountingEvent\" WHERE event_id = 'a' FOR UPDATE") + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("a",), + ("c",), + ] + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + + +@pytest.mark.asyncio +async def test_retention_makes_progress_past_pending_and_unpriced_events(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage("a", inbound=True)) + await _insert(store, _usage("b")) + await _insert(store, _usage("c", inbound=True)) + _complete_and_age(ledger_database) + _rows(ledger_database, "UPDATE \"LiteLLM_AccountingDelivery\" SET state = 'pending' WHERE event_id = 'a'") + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF, limit=1) == 1 + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("a",), + ("b",), + ] + + +@pytest.mark.asyncio +async def test_zero_cost_and_empty_delivery_sets_are_durable(writer: _Writer, ledger_database: str) -> None: + _, store = writer + await _insert(store, _usage(), targets=()) + event: Final = _cost(amount="0") + assert await _insert(store, event, targets=()) == "inserted" + assert await _insert(store, event, targets=()) == "duplicate" + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(0,)] + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + + +@pytest.mark.asyncio +async def test_retention_cannot_delete_source_while_cost_acceptance_is_uncommitted( + writer: _Writer, ledger_database: str +) -> None: + client, store = writer + await _insert(store, _usage()) + _complete_and_age(ledger_database) + inserted: Final = asyncio.Event() + release: Final = asyncio.Event() + + @asynccontextmanager + async def paused_commit() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + inserted.set() + await release.wait() + + task: Final = asyncio.create_task(_insert(AccountingLedger(paused_commit), _cost())) + try: + await asyncio.wait_for(inserted.wait(), timeout=5) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 0 + finally: + release.set() + await task + assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent" ORDER BY event_id') == [ + ("cost",), + ("usage",), + ] + + +@pytest.mark.asyncio +async def test_retention_failure_rolls_back_event_and_receipt_deletion(writer: _Writer, ledger_database: str) -> None: + client, store = writer + await _insert(store, _usage(inbound=True)) + _complete_and_age(ledger_database) + + @asynccontextmanager + async def interrupted() -> AsyncGenerator[Prisma]: + async with client.writer_db.tx() as tx: + yield tx + raise RuntimeError("injected cleanup failure") + + with pytest.raises(RuntimeError, match="cleanup failure"): + await AccountingLedger(interrupted).purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(1,)] + assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)] + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 1 + + +@pytest.mark.asyncio +async def test_ledger_reads_and_writes_use_primary_when_replica_is_configured( + writer: _Writer, ledger_database: str, monkeypatch: pytest.MonkeyPatch +) -> None: + with scratch_database("replica") as replica_url: + monkeypatch.setenv("DATABASE_URL_READ_REPLICA", replica_url) + async with connected_proxy_prisma(ledger_database) as client: + store: Final = AccountingLedger.for_client(client) + await _insert(store, _usage()) + assert await _insert(store, _usage()) == "duplicate" + assert await _insert(store, _cost()) == "inserted" + _complete_and_age(ledger_database) + assert await store.purge_completed(redelivery_before=_CUTOFF, recovery_before=_CUTOFF) == 2 + assert _rows(replica_url, "SELECT to_regclass('public.\"LiteLLM_AccountingEvent\"')") == [(None,)] + + +def _index_conditions(plan: object) -> tuple[tuple[object, object], ...]: + if isinstance(plan, dict): + own: Final = ((plan.get("Index Name"), plan.get("Index Cond")),) if "Index Name" in plan else () + children: Final = (plan.get("Plan"), *plan.get("Plans", ())) + return own + tuple(chain.from_iterable(_index_conditions(child) for child in children if child is not None)) + return tuple(chain.from_iterable(_index_conditions(item) for item in plan)) if isinstance(plan, list) else () + + +@pytest.mark.parametrize("statement", [accounting_ledger._DELETE_EVENTS, accounting_ledger._DELETE_DELIVERIES]) +def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str, statement: str) -> None: + explained: Final = "EXPLAIN (FORMAT JSON) " + statement.replace("$1", "%s") + with psycopg.connect(ledger_database) as connection: + connection.execute("SET enable_seqscan = off") + plan: Final = connection.execute(explained, ("usage",) * statement.count("$1")).fetchone() + assert plan is not None + conditions: Final = _index_conditions(plan[0]) + assert any( + name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition) + for name, condition in conditions + ), conditions diff --git a/tests/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index 57a05daca94..992c1e8e875 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -26,6 +26,20 @@ _DECISIONS_PROBE_REPLY: Final = JsonResponse( "usage": {"input_tokens": 10, "output_tokens": 1}, }, ) +_EXTRACTION_PROBE_REPLY: Final = JsonResponse( + content_type="application/json", + body={ + "result": { + "model": "clairic", + "answers": {"reachable": {"type": "text", "value": "health", "confidence": 0.9}}, + "truncated": False, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + "success": True, + "errors": [], + "messages": [], + }, +) _CONFIGURED_PROBE_REPLY: Final = JsonResponse( content_type="application/json", body={ @@ -287,6 +301,29 @@ def test_decisions_mode_health_check_resolves_the_mode_from_the_cost_map_and_sen ] +def test_extraction_mode_health_check_resolves_the_mode_from_the_cost_map_and_sends_the_extraction_probe( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + handle: Final = register_scenario(f"health-extraction-{uuid.uuid4().hex[:12]}", _EXTRACTION_PROBE_REPLY) + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + model="cloudflare/@cf/cloudflare/clairic", api_base=handle.api_base(), api_key="synthetic-cloudflare-key" + ) + report: Final = _health_report(gateway, model) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert _probes_sent_to(gateway, handle) == [ + ( + f"/{handle.scenario_id}/ai/run/@cf/cloudflare/clairic", + { + "model": "clairic", + "state": os.environ.get("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"), + "fields": {"reachable": {"type": "text", "instructions": "Any word from the text"}}, + }, + ) + ] + + @pytest.mark.parametrize("mode", ("decisions", "evaluation"), ids=("decisions", "legacy-evaluation")) def test_decisions_mode_health_check_sends_the_configured_state_and_questions(gateway: Gateway, mode: str) -> None: with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_anthropic_context_window_stop_wire.py b/tests/integration/providers/test_anthropic_context_window_stop_wire.py new file mode 100644 index 00000000000..38e3e7dffcd --- /dev/null +++ b/tests/integration/providers/test_anthropic_context_window_stop_wire.py @@ -0,0 +1,122 @@ +import json +import uuid +from collections.abc import Mapping +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +MODEL: Final = "claude-haiku-4-5" + + +def _sse(event: Mapping[str, JsonValue]) -> bytes: + return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() + + +def _message(identity: str, stop_reason: str | None) -> dict[str, JsonValue]: + return { + "id": "msg_" + identity, + "type": "message", + "role": "assistant", + "model": MODEL, + "content": [{"type": "text", "text": "alpha bravo"}], + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": 199878, "output_tokens": 12799}, + } + + +def _stream(identity: str, stop_reason: str) -> tuple[bytes, ...]: + return ( + _sse({"type": "message_start", "message": {**_message(identity, None), "content": []}}), + _sse({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _sse({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "alpha bravo"}}), + _sse({"type": "content_block_stop", "index": 0}), + _sse( + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": {"output_tokens": 12799}, + } + ), + _sse({"type": "message_stop"}), + ) + + +def _anthropic(identity: str, stop_reason: str): + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + if json.loads(request.body).get("stream"): + return Reply(content_type="text/event-stream", chunks=_stream(identity, stop_reason)) + return Reply(body=json.dumps(_message(identity, stop_reason)).encode()) + + return respond + + +def _events(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + with gateway.client.stream( + "POST", path, json={**body, "stream": True}, headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + assert response.status_code == 200, response.read().decode() + return tuple( + json.loads(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _chat_finish_reason(gateway: Gateway, model: str, *, stream: bool) -> JsonValue: + body: Final = {"model": model, "max_tokens": 64000, "messages": [{"role": "user", "content": "copy the list"}]} + if stream: + reasons: Final = tuple( + choice["finish_reason"] + for chunk in _events(gateway, "/v1/chat/completions", body) + for choice in chunk["choices"] # comprehension-ok: chunk -> choices flatten + if choice.get("finish_reason") + ) + assert len(reasons) == 1, reasons + return reasons[0] + response: Final = gateway.request("POST", "/v1/chat/completions", body) + assert response.status_code == 200, response.text + return response.json()["choices"][0]["finish_reason"] + + +def _responses_status(gateway: Gateway, model: str, *, stream: bool) -> JsonValue: + body: Final = {"model": model, "max_output_tokens": 64000, "input": "copy the list"} + if stream: + terminal: Final = tuple( + event["response"] + for event in _events(gateway, "/v1/responses", body) + if event["type"] in ("response.completed", "response.incomplete", "response.failed") + ) + assert len(terminal) == 1, terminal + return terminal[0]["status"] + response: Final = gateway.request("POST", "/v1/responses", body) + assert response.status_code == 200, response.text + return response.json()["status"] + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +@pytest.mark.parametrize( + ("stop_reason", "finish_reason"), + [("model_context_window_exceeded", "length"), ("max_tokens", "length"), ("end_turn", "stop")], +) +def test_chat_maps_anthropic_stop_reason(gateway: Gateway, stop_reason: str, finish_reason: str, stream: bool) -> None: + identity: Final = uuid.uuid4().hex + with wire_server(_anthropic(identity, stop_reason)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key="scripted-anthropic-key") + assert _chat_finish_reason(gateway, model, stream=stream) == finish_reason + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "stream"]) +@pytest.mark.parametrize( + ("stop_reason", "status"), + [("model_context_window_exceeded", "incomplete"), ("max_tokens", "incomplete"), ("end_turn", "completed")], +) +def test_responses_maps_anthropic_stop_reason(gateway: Gateway, stop_reason: str, status: str, stream: bool) -> None: + identity: Final = uuid.uuid4().hex + with wire_server(_anthropic(identity, stop_reason)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key="scripted-anthropic-key") + assert _responses_status(gateway, model, stream=stream) == status diff --git a/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py b/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py new file mode 100644 index 00000000000..7b2a48c9676 --- /dev/null +++ b/tests/integration/providers/test_cloudflare_clairic_extraction_wire.py @@ -0,0 +1,418 @@ +import json +import math +import re +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +Route = Literal["/v1/extraction", "/extraction"] + +_API_KEY: Final = "synthetic-cloudflare-key" +_CLAIRIC: Final = "cloudflare/@cf/cloudflare/clairic" +_CLAIRIC_PATH: Final = "/ai/run/@cf/cloudflare/clairic" +_INPUT_TOKENS: Final = 412 +_EMAIL: Final = "billing@acmetools.example" +_SCENARIO: Final = re.compile(r"scenario-([a-z-]+)") +_JSON_ROWS: Final = TypeAdapter(list[dict[str, JsonValue]]) +_FIELDS: Final[dict[str, JsonValue]] = { + "vendor": {"type": "text", "instructions": "Who sent the invoice?"}, + "total": {"type": "number", "unit": "USD"}, +} +_SPEND_QUERY: Final = ( + 'SELECT spend, status, call_type, model_group, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" ' + "WHERE request_id = %s" +) +_OWNED_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_REQUEST_GUARDRAIL: Final = "extraction-mask-email-in" +_RESPONSE_GUARDRAIL: Final = "extraction-mask-email-out" +_OWNED_MODEL: Final = "invoice-reader" + + +def _state(marker: str, scenario: str = "plain") -> str: + return f"Invoice from Acme Tools, total $174. scenario-{scenario} marker-{marker}" + + +def _vendor_span(state: str) -> list[JsonValue]: + start: Final = state.index("Acme Tools") + return [start, start + len("Acme Tools")] + + +def _email_span(state: str) -> list[JsonValue]: + start: Final = state.index(_EMAIL) + return [start, start + len(_EMAIL)] + + +def _answers(state: str) -> dict[str, JsonValue]: + return { + "vendor": {"type": "text", "value": "Acme Tools", "confidence": 0.97, "span": _vendor_span(state)}, + "total": {"type": "number", "value": 174, "confidence": 0.95}, + } + + +def _envelope(result: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps({"result": result, "success": True, "errors": [], "messages": []}).encode()) + + +def _clairic_reply(request: Request) -> Reply: + state: Final = JSON_OBJECT.validate_json(request.body)["state"] + assert isinstance(state, str), request.body + return _envelope( + { + "model": "clairic", + "answers": _answers(state), + "truncated": False, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + ) + + +def _expected_response(model: str, state: str) -> dict[str, JsonValue]: + return { + "model": model, + "answers": _answers(state), + "truncated": False, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + + +def _upstream_body(request: Request) -> dict[str, JsonValue]: + assert (request.method, request.target) == ("POST", _CLAIRIC_PATH), (request.method, request.target) + assert request.headers["authorization"] == f"Bearer {_API_KEY}", request.headers + return JSON_OBJECT.validate_json(request.body) + + +def _body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "state": _state(marker), "fields": _FIELDS, **extra} + + +def _extract( + gateway: Gateway, + model: str, + marker: str, + *, + route: Route = "/v1/extraction", + key: str | None = None, + **extra: JsonValue, +) -> httpx.Response: + return gateway.request("POST", route, _body(model, marker, **extra), key=key) + + +def _number(value: JsonValue) -> float: + assert isinstance(value, (int, float)) and not isinstance(value, bool), value + return float(value) + + +def _expected_spend() -> float: + cost_map: Final = JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes()) + priced: Final = cost_map.get(_CLAIRIC) + assert priced is not None, f"{_CLAIRIC} has no cost map entry" + return _INPUT_TOKENS * _number(object_value(priced)["input_cost_per_token"]) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) + return rows[0] + + +def _assert_billed(gateway: Gateway, response: httpx.Response, model: str) -> None: + expected: Final = _expected_spend() + assert math.isclose(float(response.headers["x-litellm-response-cost"]), expected, rel_tol=1e-9), response.headers + call_id: Final = response.headers["x-litellm-call-id"] + row: Final = _spend_row(call_id) + assert (row["status"], row["call_type"], row["model_group"], row["prompt_tokens"], row["completion_tokens"]) == ( + "success", + "aextraction", + model, + _INPUT_TOKENS, + 0, + ), row + assert math.isclose(_number(row["spend"]), expected, rel_tol=1e-9), row + listed: Final = gateway.request("GET", "/spend/logs", params={"request_id": call_id}) + assert listed.status_code == 200, listed.text + (reported,) = _JSON_ROWS.validate_json(listed.content) + assert (reported["request_id"], reported["model_group"]) == (call_id, model), reported + assert math.isclose(_number(reported["spend"]), expected, rel_tol=1e-9), reported + + +@pytest.mark.parametrize("route", ("/v1/extraction", "/extraction")) +def test_extraction_reaches_clairic_on_both_routes_and_is_billed(gateway: Gateway, route: Route) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = _extract(gateway, model, marker, route=route, metadata={"tags": [f"audit-{marker}"]}) + assert response.status_code == 200, response.text + assert response.json() == _expected_response(model, _state(marker)) + (request,) = wire.drain() + assert _upstream_body(request) == {"model": "clairic", "state": _state(marker), "fields": _FIELDS} + _assert_billed(gateway, response, model) + + +_INVALID_BODIES: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ("unknown_type", {"fields": {"vendor": {"type": "paragraph"}}}), + ("list_without_items", {"fields": {"items": {"type": "list"}}}), + ("remote_image", {"images": ["https://images.example/receipt.png"]}), +) + + +@pytest.mark.parametrize(("label", "change"), _INVALID_BODIES, ids=[label for label, _ in _INVALID_BODIES]) +def test_invalid_extraction_bodies_are_refused_before_any_upstream_call( + gateway: Gateway, label: str, change: dict[str, JsonValue] +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request("POST", "/v1/extraction", {**_body(model, marker), **change}) + assert response.status_code == 400, (label, response.text) + assert "Invalid Extraction request" in response.text, (label, response.text) + assert wire.drain() == (), label + + +@pytest.mark.parametrize(("param", "value"), (("temperature", 0.2), ("stream", True)), ids=("temperature", "stream")) +def test_unsupported_top_level_params_are_refused_unless_the_deployment_drops_params( + gateway: Gateway, param: str, value: JsonValue +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + strict: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + dropping: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY, drop_params=True) + refused: Final = _extract(gateway, strict, marker, **{param: value}) + assert refused.status_code == 400, refused.text + assert param in refused.text, refused.text + assert wire.drain() == () + dropped: Final = _extract(gateway, dropping, marker, **{param: value}) + assert dropped.status_code == 200, dropped.text + assert dropped.headers["content-type"].startswith("application/json"), dropped.headers + assert dropped.json() == _expected_response(dropping, _state(marker)) + (request,) = wire.drain() + assert _upstream_body(request) == {"model": "clairic", "state": _state(marker), "fields": _FIELDS} + _assert_billed(gateway, dropped, dropping) + + +def test_key_checks_on_extraction_match_chat(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + anonymous: Final = gateway.client.post("/v1/extraction", json=_body(model, marker)) + assert anonymous.status_code == 401, anonymous.text + restricted: Final = scenario.key(models=[f"other-{uuid.uuid4().hex}"]) + refused: Final = _extract(gateway, model, marker, key=restricted) + chat_refused: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}]}, + key=restricted, + ) + assert 400 <= refused.status_code < 500, refused.text + assert refused.status_code == chat_refused.status_code, (refused.text, chat_refused.text) + assert wire.drain() == () + allowed: Final = scenario.key(models=[model], max_budget=1e-06) + first: Final = _extract(gateway, model, marker, key=allowed) + assert first.status_code == 200, first.text + assert len(wire.drain()) == 1 + blocked, upstream_calls = eventually( + lambda: (_extract(gateway, model, uuid.uuid4().hex, key=allowed), len(wire.drain())), + lambda answer: answer[0].status_code != 200, + seconds=70, + ) + chat_blocked: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}]}, + key=allowed, + ) + assert 400 <= blocked.status_code < 500, blocked.text + assert blocked.status_code == chat_blocked.status_code, (blocked.text, chat_blocked.text) + assert upstream_calls == 0 + assert wire.drain() == () + + +def test_a_small_key_tpm_limit_serves_extraction_and_counts_its_input_tokens(gateway: Gateway) -> None: + with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + key: Final = scenario.key(models=[model], tpm_limit=100) + first: Final = _extract(gateway, model, uuid.uuid4().hex, key=key) + assert first.status_code == 200, first.text + assert len(wire.drain()) == 1 + limited, upstream_calls = eventually( + lambda: (_extract(gateway, model, uuid.uuid4().hex, key=key), len(wire.drain())), + lambda answer: answer[0].status_code != 200, + seconds=70, + ) + assert limited.status_code == 429, limited.text + assert upstream_calls == 0 + + +def test_cloudflare_errors_keep_their_status_and_log_an_unbilled_failure(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + refusal: Final = json.dumps( + { + "result": None, + "success": False, + "errors": [{"code": 5006, "message": f"min above max {marker}"}], + "messages": [], + } + ).encode() + with wire_server(lambda _: Reply(status=422, body=refusal)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY) + response: Final = _extract(gateway, model, marker) + assert response.status_code == 422, response.text + assert f"min above max {marker}" in response.text, response.text + assert len(wire.drain()) >= 1 + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == ( + "failure", + "aextraction", + model, + 0.0, + ), row + + +def _owned_reply(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + state: Final = body["state"] + assert isinstance(state, str), request.body + scenario: Final = _SCENARIO.search(state) + assert scenario is not None, state + if scenario.group(1) != "mask-out": + return _clairic_reply(request) + return _envelope( + { + "model": "clairic", + "answers": { + "vendor": {"type": "text", "value": "Acme Tools", "confidence": 0.97, "span": _vendor_span(state)}, + "contact_email": {"type": "text", "value": _EMAIL, "confidence": 0.9, "span": _email_span(state)}, + }, + "truncated": False, + "state_text": state, + "usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0}, + } + ) + + +def _owned_config(upstream_url: str) -> dict[str, JsonValue]: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + email_filter: Final[list[JsonValue]] = [{"pattern_type": "prebuilt", "pattern_name": "email", "action": "MASK"}] + return { + **base, + "guardrails": [ + { + "guardrail_name": _REQUEST_GUARDRAIL, + "litellm_params": {"guardrail": "litellm_content_filter", "mode": "pre_call", "patterns": email_filter}, + }, + { + "guardrail_name": _RESPONSE_GUARDRAIL, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "patterns": email_filter, + }, + }, + ], + "model_list": [ + { + "model_name": _OWNED_MODEL, + "litellm_params": {"model": _CLAIRIC, "api_base": upstream_url, "api_key": _API_KEY}, + } + ], + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: OwnedProxy + upstream: Wire + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with wire_server(_owned_reply) as upstream: + directory: Final = tmp_path_factory.mktemp("clairic-extraction-owned") + config_path: Final = directory / "clairic-extraction.yaml" + config_path.write_text(yaml.safe_dump(_owned_config(upstream.url))) + with owned_proxy_process( + session_gateway, + directory, + {}, + config=config_path, + workers=2, + extra_arguments=("--max_tokens", "256", "--temperature", "0.2"), + healthcheck=WorkerHealthcheck.LENIENT, + ) as proxy: + yield Rig(proxy, upstream) + + +def _owned_extract(rig: Rig, state: str, **extra: JsonValue) -> httpx.Response: + return rig.proxy.gateway.request( + "POST", "/v1/extraction", {"model": _OWNED_MODEL, "state": state, "fields": _FIELDS, **extra} + ) + + +def _applied_guardrails(response: httpx.Response) -> frozenset[str]: + return frozenset(name.strip() for name in response.headers.get("x-litellm-applied-guardrails", "").split(",")) + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_cli_generation_defaults_never_reach_an_extraction_call(rig: Rig) -> None: + state: Final = _state(uuid.uuid4().hex) + response: Final = _owned_extract(rig, state) + assert response.status_code == 200, response.text + assert response.json() == _expected_response(_OWNED_MODEL, state) + (request,) = rig.upstream.drain() + assert _upstream_body(request) == {"model": "clairic", "state": state, "fields": _FIELDS} + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_a_request_guardrail_masks_the_state_before_clairic_and_drops_spans_into_it(rig: Rig) -> None: + state: Final = f"{_state(uuid.uuid4().hex, 'mask-in')} Questions go to {_EMAIL}." + response: Final = _owned_extract(rig, state, guardrails=[_REQUEST_GUARDRAIL]) + assert response.status_code == 200, response.text + assert _REQUEST_GUARDRAIL in _applied_guardrails(response), response.headers + (request,) = rig.upstream.drain() + sent: Final = _upstream_body(request) + sent_state: Final = sent["state"] + assert isinstance(sent_state, str) and _EMAIL not in sent_state, sent + assert sent_state.startswith(state.removesuffix(f"{_EMAIL}.")) and "REDACTED" in sent_state, sent + assert sent == {"model": "clairic", "state": sent_state, "fields": _FIELDS} + answers: Final = object_value(response.json()["answers"]) + assert object_value(answers["vendor"]) == {"type": "text", "value": "Acme Tools", "confidence": 0.97}, answers + + +@pytest.mark.timeout(_OWNED_TIMEOUT) +def test_a_response_guardrail_masks_text_answers_and_state_text(rig: Rig) -> None: + state: Final = f"{_state(uuid.uuid4().hex, 'mask-out')} Questions go to {_EMAIL}." + fields: Final[dict[str, JsonValue]] = {**_FIELDS, "contact_email": {"type": "text"}} + response: Final = rig.proxy.gateway.request( + "POST", + "/v1/extraction", + { + "model": _OWNED_MODEL, + "state": state, + "fields": fields, + "guardrails": [_RESPONSE_GUARDRAIL], + "options": {"return_state_text": True}, + }, + ) + assert response.status_code == 200, response.text + assert _RESPONSE_GUARDRAIL in _applied_guardrails(response), response.headers + (request,) = rig.upstream.drain() + assert _upstream_body(request)["state"] == state + body: Final = response.json() + state_text: Final = body["state_text"] + assert isinstance(state_text, str) and _EMAIL not in state_text and "REDACTED" in state_text, body + answers: Final = object_value(body["answers"]) + contact: Final = object_value(answers["contact_email"]) + assert "span" not in contact, answers + masked: Final = contact["value"] + assert isinstance(masked, str) and _EMAIL not in masked and "REDACTED" in masked, answers + assert object_value(answers["vendor"]) == {"type": "text", "value": "Acme Tools", "confidence": 0.97}, answers diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 04b78985b20..5af1532e9ed 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -1,7 +1,7 @@ import json import math import uuid -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import Final @@ -140,6 +140,38 @@ _PROVIDERS: Final = ( "databricks/databricks-openjev-qwen35-4b", ), ) +_LOCAL_PROVIDERS: Final = ( + pytest.param( + _Provider("ollama", "ollama/d3:4b", "/v1/systemone", "d3:4b", None, False, None), id="regression-ollama" + ), + pytest.param( + _Provider("ollama_chat", "ollama_chat/d3:4b", "/v1/systemone", "d3:4b", None, False, None), + id="regression-ollama_chat", + ), + pytest.param( + _Provider("hosted_vllm", "hosted_vllm/d3:4b", "/v1/systemone", "d3:4b", None, False, None), + id="guard-hosted_vllm", + ), + pytest.param( + _Provider("lm_studio", "lm_studio/d3:4b", "/v1/systemone", "d3:4b", None, False, None), + id="guard-lm_studio", + ), +) +_LOCAL_STATE: Final = "I was charged twice for order 4417." +_LOCAL_QUESTIONS: Final[Mapping[str, JsonValue]] = { + "billing": {"type": "noul", "instructions": "Is this a billing issue?"} +} +_LOCAL_ANSWERS: Final[Mapping[str, JsonValue]] = {"billing": {"type": "noul", "noul": 0.91}} +_LOCAL_OPENAI_QUESTIONS: Final[Sequence[JsonValue]] = ( + {"type": "predicate", "name": "billing", "instructions": "Is this a billing issue?"}, +) +_LOCAL_OPENAI_ANSWERS: Final[Sequence[JsonValue]] = ({"type": "predicate", "name": "billing", "probability": 0.91},) +_LOCAL_OPENAI_USAGE: Final[Mapping[str, JsonValue]] = { + **_USAGE, + "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 370, +} _PERPLEXITY: Final = _PROVIDERS[0] _OPENROUTER: Final = next(provider for provider in _PROVIDERS if provider.name == "openrouter") _OPENROUTER_CHAT_MODEL: Final = "openrouter/openai/gpt-5-mini" @@ -230,6 +262,26 @@ def _upstream_calls(gateway: Gateway, handle: ScenarioHandle) -> list[dict[str, return _calls_to(_observed_requests(gateway), handle) +def _sdk_inference_call(gateway: Gateway, handle: ScenarioHandle, provider: _Provider) -> dict[str, JsonValue]: + observed: Final = _upstream_calls(gateway, handle) + inference_path: Final = f"/{handle.scenario_id}/v1/systemone" + (call,) = tuple(request for request in observed if request["path"] == inference_path) + metadata_calls: Final = tuple(request for request in observed if request["path"] != inference_path) + assert metadata_calls in ( + (), + ( + { + "path": f"{inference_path}/api/show", + "authorization": "", + "body": {"name": provider.body_model}, + "method": "POST", + "api_key": "", + }, + ), + ), observed + return call + + def _spend_row(call_id: str) -> dict[str, JsonValue]: rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70) return rows[0] @@ -345,6 +397,174 @@ def test_each_provider_gets_its_own_path_key_and_body_and_is_billed(gateway: Gat assert math.isclose(_number(row["spend"]), expected_spend, rel_tol=1e-9), row +@pytest.fixture(autouse=True) +def _without_ambient_ollama_settings(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + + +@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS) +@pytest.mark.parametrize("route", ("/v1/systemone", "/v1/decisions"), ids=("systemone", "decisions")) +@pytest.mark.parametrize( + ("base_suffix", "api_key"), + ( + pytest.param("", None, id="root-no-key"), + pytest.param("/", _API_KEY, id="root-slash-key"), + pytest.param("/v1", None, id="v1-no-key"), + pytest.param("/v1/", _API_KEY, id="v1-slash-key"), + ), +) +def test_local_decisions_routes_preserve_wire_auth_usage_and_spend( + gateway: Gateway, provider: _Provider, route: str, base_suffix: str, api_key: str | None +) -> None: + native_reply: Final[Mapping[str, JsonValue]] = { + "model": provider.body_model, + "answers": dict(_LOCAL_ANSWERS), + "usage": _USAGE, + } + request_body: Final[Mapping[str, JsonValue]] = ( + {"state": _LOCAL_STATE, "questions": dict(_LOCAL_QUESTIONS)} + if route == "/v1/systemone" + else {"input": _LOCAL_STATE, "questions": list(_LOCAL_OPENAI_QUESTIONS)} + ) + expected_reply: Final = ( + native_reply + if route == "/v1/systemone" + else {"model": provider.body_model, "answers": list(_LOCAL_OPENAI_ANSWERS), "usage": dict(_LOCAL_OPENAI_USAGE)} + ) + with gateway.scenario() as scenario: + handle: Final = _register(scenario, dict(native_reply)) + model: Final = scenario.model( + model=provider.model, api_base=f"{handle.api_base()}{base_suffix}", api_key=api_key + ) + response: Final = gateway.request("POST", route, {"model": model, **request_body}) + assert response.status_code == 200, response.text + assert response.json() == expected_reply + assert response.headers["x-litellm-model-group"] == model + assert float(response.headers.get("x-litellm-response-cost", "0")) == 0.0 + (call,) = _upstream_calls(gateway, handle) + assert (call["path"], call["authorization"], call["body"]) == ( + f"/{handle.scenario_id}/v1/systemone", + f"Bearer {api_key}" if api_key else "", + {"model": provider.body_model, "state": _LOCAL_STATE, "questions": _LOCAL_QUESTIONS}, + ) + row: Final = _spend_row(response.headers["x-litellm-call-id"]) + assert {key: value for key, value in row.items() if key != "request_tags"} == { + "spend": 0.0, + "status": "success", + "call_type": "asystemone" if route == "/v1/systemone" else "adecisions", + "model_group": model, + "custom_llm_provider": provider.name, + "api_base": f"{handle.api_base()}/v1/systemone", + "prompt_tokens": _USAGE["input_tokens"], + "completion_tokens": _USAGE["output_tokens"], + } + + +@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS) +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +async def test_local_systemone_sdk_preserves_native_wire( + gateway: Gateway, provider: _Provider, asynchronous: bool +) -> None: + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _answer_body(provider)) + response: Final = ( + await litellm.asystemone( + model=provider.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=f"{handle.api_base()}/v1/" + ) + if asynchronous + else litellm.systemone( + model=provider.model, state=_STATE, questions=_SDK_QUESTIONS, api_base=handle.api_base() + ) + ) + assert response.model_dump(mode="json") == _answer_body(provider) + call: Final = _sdk_inference_call(gateway, handle, provider) + assert (call["path"], call["authorization"], call["body"]) == ( + f"/{handle.scenario_id}/v1/systemone", + "", + {"model": provider.body_model, "state": _STATE, "questions": _QUESTIONS}, + ) + + +@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS) +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +async def test_local_decisions_sdk_translates_openai_format( + gateway: Gateway, provider: _Provider, asynchronous: bool +) -> None: + questions: Final = TypeAdapter(list[dict[str, object]]).validate_python(_LOCAL_OPENAI_QUESTIONS) + native_reply: Final[Mapping[str, JsonValue]] = { + "model": provider.body_model, + "answers": dict(_LOCAL_ANSWERS), + "usage": _USAGE, + } + with gateway.scenario() as scenario: + handle: Final = _register(scenario, dict(native_reply)) + response: Final = ( + await litellm.adecisions( + model=provider.model, input=_LOCAL_STATE, questions=questions, api_base=f"{handle.api_base()}/v1/" + ) + if asynchronous + else litellm.decisions( + model=provider.model, input=_LOCAL_STATE, questions=questions, api_base=handle.api_base() + ) + ) + assert response.model_dump(mode="json") == { + "model": provider.body_model, + "answers": list(_LOCAL_OPENAI_ANSWERS), + "usage": dict(_LOCAL_OPENAI_USAGE), + } + call: Final = _sdk_inference_call(gateway, handle, provider) + assert (call["path"], call["authorization"], call["body"]) == ( + f"/{handle.scenario_id}/v1/systemone", + "", + {"model": provider.body_model, "state": _LOCAL_STATE, "questions": _LOCAL_QUESTIONS}, + ) + + +@pytest.mark.parametrize("provider", _LOCAL_PROVIDERS[:2]) +@pytest.mark.parametrize( + ("environment_base", "api_key", "environment_key", "authorization"), + ( + pytest.param(True, None, _ENV_KEY, f"Bearer {_ENV_KEY}", id="environment-base-and-key"), + pytest.param(False, _API_KEY, _ENV_KEY, f"Bearer {_API_KEY}", id="explicit-overrides-environment"), + pytest.param(False, "fake-api-key", _ENV_KEY, "", id="fake-key-suppresses-environment"), + pytest.param(False, None, None, "", id="unrelated-keys-not-forwarded"), + ), +) +def test_ollama_decisions_sdk_respects_environment_and_explicit_auth( + gateway: Gateway, + monkeypatch: pytest.MonkeyPatch, + provider: _Provider, + environment_base: bool, + api_key: str | None, + environment_key: str | None, + authorization: str, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "synthetic-unrelated-openai-key") + monkeypatch.setenv("LITELLM_API_KEY", "synthetic-unrelated-litellm-key") + if environment_key is not None: + monkeypatch.setenv("OLLAMA_API_KEY", environment_key) + with gateway.scenario() as scenario: + handle: Final = _register(scenario, _answer_body(provider)) + unused: Final = _register(scenario, _answer_body(provider)) + monkeypatch.setenv("OLLAMA_API_BASE", handle.api_base() if environment_base else unused.api_base()) + response: Final = litellm.systemone( + model=provider.model, + state=_STATE, + questions=_SDK_QUESTIONS, + api_base=None if environment_base else f"{handle.api_base()}/v1/", + api_key=api_key, + ) + assert response.model_dump(mode="json") == _answer_body(provider) + call: Final = _sdk_inference_call(gateway, handle, provider) + assert (call["path"], call["authorization"], call["body"]) == ( + f"/{handle.scenario_id}/v1/systemone", + authorization, + {"model": provider.body_model, "state": _STATE, "questions": _QUESTIONS}, + ) + assert _upstream_calls(gateway, unused) == [] + + def test_test_connection_evaluation_mode_uses_typesafe_decisions_path(gateway: Gateway) -> None: provider: Final = _PROVIDERS[1] with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_gemini_agents_wire.py b/tests/integration/providers/test_gemini_agents_wire.py index a619132d067..fd8f6b6ca63 100644 --- a/tests/integration/providers/test_gemini_agents_wire.py +++ b/tests/integration/providers/test_gemini_agents_wire.py @@ -75,9 +75,9 @@ def test_agents_forwards_caller_supplied_key_to_chosen_api_base(gateway: Gateway assert request.headers["x-goog-api-key"] == "synthetic-caller-key" assert _ENV_GEMINI_KEY not in request.headers.values() assert json.loads(request.body) == { - "name": _AGENT_NAME, + "id": _AGENT_NAME, "base_agent": "waverunner", - "instructions": "make slides", + "system_instruction": "make slides", } return Reply(body=json.dumps(_AGENT_REPLY).encode()) diff --git a/tests/integration/providers/test_gemini_grounding_annotations_wire.py b/tests/integration/providers/test_gemini_grounding_annotations_wire.py new file mode 100644 index 00000000000..11df7c214f3 --- /dev/null +++ b/tests/integration/providers/test_gemini_grounding_annotations_wire.py @@ -0,0 +1,65 @@ +import json +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + + +def _grounded_reply(request: Request) -> Reply: + body: Final = { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "The passage"}]}, + "finishReason": "STOP", + "groundingMetadata": { + "groundingChunks": [ + {"retrievedContext": {"uri": "gs://documents/source"}}, + {"web": {"uri": "https://example.com/first", "title": "First"}}, + {"web": {"uri": "https://example.com/second", "title": "Second"}}, + ], + "groundingSupports": [ + {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": [0, 99, 2, 1]} + ], + }, + } + ], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 2, "totalTokenCount": 7}, + } + assert json.loads(request.body)["contents"] == [{"role": "user", "parts": [{"text": "Find sources"}]}] + if ":streamGenerateContent" in request.target: + return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(body)}\n\n".encode(),)) + return Reply(body=json.dumps(body).encode()) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +def test_chat_grounding_annotations_include_every_web_source(gateway: Gateway, stream: bool) -> None: + with wire_server(_grounded_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="gemini/gemini-3.8-flash", api_base=wire.url, api_key="synthetic-key") + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "Find sources"}], "stream": stream}, + ) + assert response.status_code == 200, response.text + payloads: Final = ( + [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] + if stream + else [response.json()] + ) + messages: Final = tuple(payload["choices"][0]["delta" if stream else "message"] for payload in payloads) + annotations: Final = tuple(message["annotations"] for message in messages if message.get("annotations")) + assert annotations == ( + [ + { + "type": "url_citation", + "url_citation": { + "start_index": 4, + "end_index": 11, + "url": f"https://example.com/{name.lower()}", + "title": name, + }, + } + for name in ("Second", "First") + ], + ), response.text diff --git a/tests/integration/routing/_reasoning_effort_levels_support.py b/tests/integration/routing/_reasoning_effort_levels_support.py index 8d43695e419..d295f6e6d56 100644 --- a/tests/integration/routing/_reasoning_effort_levels_support.py +++ b/tests/integration/routing/_reasoning_effort_levels_support.py @@ -32,6 +32,7 @@ GEMINI: Final = "gemini-3-pro-preview" ZAI_GLM: Final = "glm-5" ECHOED_KEY_MODEL: Final = f"{KIMI}-echoed-key" BEDROCK_CLAUDE: Final = f"us.anthropic.{CLAUDE}" +DATABRICKS_CLAUDE: Final = f"databricks-{CLAUDE}" AZURE_API_VERSION: Final = "2025-04-01-preview" SYNTHETIC_KEY: Final = "synthetic-provider-key" _MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") @@ -215,6 +216,10 @@ def backend(kind: str, wire_url: str) -> Backend: f"bedrock/invoke/{BEDROCK_CLAUDE}", MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY, "aws_region_name": "us-east-1"}), ) + case "databricks_claude": + return Backend( + f"databricks/{DATABRICKS_CLAUDE}", MappingProxyType({"api_base": wire_url, "api_key": SYNTHETIC_KEY}) + ) case _: raise AssertionError(f"unknown backend kind {kind}") diff --git a/tests/integration/routing/test_enforce_model_rate_limits_daily_limits.py b/tests/integration/routing/test_enforce_model_rate_limits_daily_limits.py new file mode 100644 index 00000000000..c670171a149 --- /dev/null +++ b/tests/integration/routing/test_enforce_model_rate_limits_daily_limits.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import gateway_from_environment +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import owned_config +from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GROUP: Final = "daily-limited-group" +LIMITED_ID: Final = "free-tier-two-a-day" +UNLIMITED_ID: Final = "paid-tier" +RPD: Final = 2 +EXTRA_REQUESTS: Final = 2 + + +@dataclass(frozen=True, slots=True) +class Rig: + owned: OwnedProxy + limited: Wire + unlimited: Wire + + +def _reply(request: Request) -> Reply: + if is_model_info_probe(request): + return model_list_reply() + return chat_reply("chatcmpl-daily", OPENAI_MODEL, "served", stream=False) + + +def _deployment(model_id: str, wire: Wire, **params: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": GROUP, + "litellm_params": { + "model": OPENAI_MODEL, + "api_base": wire.url + "/v1", + "api_key": "synthetic-openai-key", + **params, + }, + "model_info": {"id": model_id}, + } + + +def _config(directory: Path, limited: Wire, unlimited: Wire) -> Path: + return owned_config( + directory, + [_deployment(LIMITED_ID, limited, rpd=RPD, order=1), _deployment(UNLIMITED_ID, unlimited, order=2)], + litellm_settings={"cache": False}, + router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 1}, + ) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("daily-limits") + with ( + gateway_from_environment() as gateway, + wire_server(_reply) as limited, + wire_server(_reply) as unlimited, + owned_proxy_process(gateway, directory, {}, config=_config(directory, limited, unlimited)) as owned, + ): + yield Rig(owned, limited, unlimited) + + +def _chat(rig: Rig) -> httpx.Response: + return rig.owned.gateway.request( + "POST", + "/v1/chat/completions", + {"model": GROUP, "messages": [{"role": "user", "content": f"daily {uuid.uuid4().hex}"}]}, + ) + + +def _forwarded(wire: Wire) -> int: + return sum(1 for request in wire.drain() if not is_model_info_probe(request)) + + +def test_requests_past_the_deployment_rpd_are_served_by_the_other_deployment_in_the_group(rig: Rig) -> None: + responses: Final = tuple(_chat(rig) for _ in range(RPD + EXTRA_REQUESTS)) + + assert [response.status_code for response in responses] == [200] * (RPD + EXTRA_REQUESTS), [ + response.text for response in responses + ] + assert [response.headers["x-litellm-model-id"] for response in responses] == [LIMITED_ID] * RPD + [ + UNLIMITED_ID + ] * EXTRA_REQUESTS + assert (_forwarded(rig.limited), _forwarded(rig.unlimited)) == (RPD, EXTRA_REQUESTS) diff --git a/tests/integration/routing/test_reasoning_effort_levels_wire.py b/tests/integration/routing/test_reasoning_effort_levels_wire.py index 3d88aad547c..d15867ac572 100644 --- a/tests/integration/routing/test_reasoning_effort_levels_wire.py +++ b/tests/integration/routing/test_reasoning_effort_levels_wire.py @@ -1,4 +1,5 @@ import asyncio +import json from collections.abc import Callable, Mapping from dataclasses import dataclass from types import MappingProxyType @@ -8,9 +9,11 @@ import anthropic import httpx import openai import pytest +from integration._support.anthropic_sse import message_json from integration._support.client import Gateway, eventually, list_value, object_value, string_value from integration._support.openai_wire import openai_error -from integration._support.wire import Request, wire_server +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server from integration.routing._reasoning_effort_levels_support import ( AZURE_API_VERSION, CLAUDE, @@ -34,6 +37,7 @@ from integration.routing._reasoning_effort_levels_support import ( marker_of, messages_request, new_marker, + path_of, prompt, provider_reply, responses_request, @@ -41,6 +45,8 @@ from integration.routing._reasoning_effort_levels_support import ( ) from pydantic import JsonValue +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + CHAT: Final = "/v1/chat/completions" MESSAGES: Final = "/v1/messages" RESPONSES: Final = "/v1/responses" @@ -49,6 +55,7 @@ _ADAPTIVE: Final[dict[str, JsonValue]] = {"type": "adaptive"} _STRING_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": "2048"} _XHIGH_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 9000} _BOOL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": True} +_FALSE_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": False} _NULL_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": None} _LIST_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": [2048]} _EMPTY_BUDGET: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": ""} @@ -67,6 +74,8 @@ _KIMI_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort _CLAUDE_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(CLAUDE_LEVELS)}) _GPT_INFO: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": list(GPT_LEVELS)}) _UNDECLARED: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +_EMPTY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": []}) +_NONE_ONLY_DECLARATION: Final[Mapping[str, JsonValue]] = MappingProxyType({"reasoning_effort_levels": ["none"]}) _ECHOED_MAP_KEY: Final[Mapping[str, JsonValue]] = MappingProxyType( {"key": "azure/chat-latest", "litellm_provider": "azure"} ) @@ -134,7 +143,7 @@ def _invoke_chat(effort: str, thinking: Mapping[str, JsonValue], betas: tuple[st "messages": [{"role": "user", "content": [{"type": "text", "text": prompt(marker)}]}], "thinking": dict(thinking), "output_config": {"effort": effort}, - "max_tokens": 4096, + "max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "anthropic_version": "bedrock-2023-05-31", "anthropic_beta": list(betas), } @@ -144,6 +153,14 @@ def _with_output_config(expected: Expected, output_config: JsonValue) -> Expecte return lambda marker: {**expected(marker), "output_config": output_config} +def _with_fields(expected: Expected, fields: Mapping[str, JsonValue]) -> Expected: + return lambda marker: {**expected(marker), **fields} + + +def _without(expected: Expected, field: str) -> Expected: + return lambda marker: {key: value for key, value in expected(marker).items() if key != field} + + def _converse_chat(effort: str, thinking: Mapping[str, JsonValue]) -> Expected: return lambda marker: { "messages": [{"role": "user", "content": [{"text": prompt(marker)}]}], @@ -320,6 +337,7 @@ _OPENAI_RESPONSES: Final = "/v1/responses" _AZURE_RESPONSES: Final = f"/openai/responses?api-version={AZURE_API_VERSION}" _UNSUPPORTED: Final = "does not support parameters" _GEMINI: Final = f"/models/{GEMINI}:generateContent" +_DATABRICKS: Final = "/chat/completions" CELLS: Final[Mapping[str, Cell]] = MappingProxyType( { @@ -1750,6 +1768,139 @@ CELLS: Final[Mapping[str, Cell]] = MappingProxyType( _anthropic_chat_without_effort(_ADAPTIVE), _ANTHROPIC, ), + "chat-anthropic-false-budget-400": Cell( + CHAT, "anthropic", _CLAUDE_INFO, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET + ), + "chat-anthropic-undeclared-null-budget-400": Cell( + CHAT, "anthropic", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET + ), + "chat-anthropic-undeclared-false-budget-drop-params-sends-adaptive-without-effort": Cell( + CHAT, + "anthropic", + _UNDECLARED, + {"thinking": _FALSE_BUDGET, "drop_params": True}, + _anthropic_chat_without_effort(_ADAPTIVE), + _ANTHROPIC, + ), + "chat-anthropic-undeclared-missing-budget-sends-low": Cell( + CHAT, + "anthropic", + _UNDECLARED, + {"thinking": _MISSING_BUDGET}, + _anthropic_chat("low", _ADAPTIVE), + _ANTHROPIC, + ), + "messages-anthropic-undeclared-false-budget-400": Cell( + MESSAGES, "anthropic", _UNDECLARED, {"thinking": _FALSE_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET + ), + "messages-anthropic-empty-string-budget-400": Cell( + MESSAGES, "anthropic", _CLAUDE_INFO, {"thinking": _EMPTY_BUDGET}, None, _ANTHROPIC, 400, _NON_NUMERIC_BUDGET + ), + "messages-anthropic-undeclared-null-budget-drop-params-sends-adaptive-without-effort": Cell( + MESSAGES, + "anthropic", + _UNDECLARED, + {"thinking": _NULL_BUDGET, "drop_params": True}, + _native_messages(None, _ADAPTIVE), + _ANTHROPIC, + ), + "responses-anthropic-false-budget-400": Cell( + RESPONSES, + "anthropic", + _CLAUDE_INFO, + {"thinking": _FALSE_BUDGET}, + None, + _ANTHROPIC, + 400, + _NON_NUMERIC_BUDGET, + ), + "chat-converse-undeclared-false-budget-400": Cell( + CHAT, + "bedrock_converse", + _UNDECLARED, + {"thinking": _FALSE_BUDGET}, + None, + _CONVERSE, + 400, + _NON_NUMERIC_BUDGET, + ), + "messages-invoke-undeclared-null-budget-400": Cell( + MESSAGES, "bedrock_invoke", _UNDECLARED, {"thinking": _NULL_BUDGET}, None, _INVOKE, 400, _NON_NUMERIC_BUDGET + ), + "chat-databricks-claude-undeclared-false-budget-400": Cell( + CHAT, + "databricks_claude", + _UNDECLARED, + {"thinking": _FALSE_BUDGET}, + None, + _DATABRICKS, + 400, + _NON_NUMERIC_BUDGET, + ), + "chat-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + CHAT, + "anthropic", + _EMPTY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _anthropic_chat_without_effort(_ADAPTIVE), + _ANTHROPIC, + ), + "chat-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + CHAT, + "anthropic", + _NONE_ONLY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _anthropic_chat_without_effort(_ADAPTIVE), + _ANTHROPIC, + ), + "chat-anthropic-empty-declaration-budget-keeps-its-effort": Cell( + CHAT, + "anthropic", + _EMPTY_DECLARATION, + {"thinking": _BUDGET}, + _anthropic_chat("medium", _ADAPTIVE), + _ANTHROPIC, + ), + "chat-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + CHAT, + "bedrock_converse", + _EMPTY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _converse_chat_without_effort(_ADAPTIVE), + _CONVERSE, + ), + "messages-anthropic-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + MESSAGES, + "anthropic", + _EMPTY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _native_messages(None, _ADAPTIVE), + _ANTHROPIC, + ), + "messages-anthropic-none-only-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + MESSAGES, + "anthropic", + _NONE_ONLY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _native_messages(None, _ADAPTIVE), + _ANTHROPIC, + ), + "messages-anthropic-empty-declaration-budget-keeps-its-effort": Cell( + MESSAGES, + "anthropic", + _EMPTY_DECLARATION, + {"thinking": _BUDGET}, + _native_messages({"effort": "medium"}, _ADAPTIVE), + _ANTHROPIC, + ), + "messages-converse-empty-declaration-budget-drop-params-sends-adaptive-without-effort": Cell( + MESSAGES, + "bedrock_converse", + _EMPTY_DECLARATION, + {"thinking": _BUDGET, "drop_params": True}, + _converse_messages(None), + _CONVERSE, + ), } ) @@ -1815,20 +1966,65 @@ def _bedrock_native_schema(body: Mapping[str, JsonValue], *, invoke: bool) -> Js return JSON_OBJECT.validate_json(string_value(json_schema["schema"])) +_CONVERSE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = { + "textFormat": { + "type": "json_schema", + "structure": {"jsonSchema": {"name": "answer", "schema": json.dumps(_NATIVE_OUTPUT_SCHEMA)}}, + } +} +_INVOKE_NATIVE_OUTPUT: Final[dict[str, JsonValue]] = { + "effort": "high", + "format": {"type": "json_schema", "schema": _NATIVE_OUTPUT_SCHEMA}, +} +_JSON_TOOL: Final = "json_tool_call" +_CONVERSE_TOOL_SCHEMA: Final[dict[str, JsonValue]] = { + key: value for key, value in _NATIVE_OUTPUT_SCHEMA.items() if key != "additionalProperties" +} +_CONVERSE_WITH_NATIVE_OUTPUT: Final = _with_fields( + _converse_chat("high", _ADAPTIVE), {"outputConfig": _CONVERSE_NATIVE_OUTPUT} +) +_INVOKE_WITH_NATIVE_OUTPUT: Final = _with_output_config( + _invoke_chat("high", _ADAPTIVE, ("structured-outputs-2025-11-13",)), _INVOKE_NATIVE_OUTPUT +) +_CONVERSE_WITH_JSON_TOOL: Final = _with_fields( + _converse_chat("high", _ADAPTIVE), + { + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": _JSON_TOOL, + "description": _JSON_TOOL, + "inputSchema": {"json": _CONVERSE_TOOL_SCHEMA}, + } + } + ] + } + }, +) +_INVOKE_WITH_JSON_TOOL: Final = _with_fields( + _without(_invoke_chat("high", _ADAPTIVE, ()), "anthropic_beta"), + {"tools": [{"name": _JSON_TOOL, "input_schema": _NATIVE_OUTPUT_SCHEMA}]}, +) + + @pytest.mark.parametrize( - ("kind", "target", "fields"), + ("kind", "target", "fields", "expected"), [ - pytest.param("bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, id="converse-budget"), + pytest.param( + "bedrock_converse", _CONVERSE, {"thinking": _BUDGET}, _CONVERSE_WITH_NATIVE_OUTPUT, id="converse-budget" + ), pytest.param( "bedrock_invoke", _INVOKE, {"thinking": _BUDGET, "output_config": {"effort": "medium"}}, + _INVOKE_WITH_NATIVE_OUTPUT, id="invoke-output-config-beside-budget", ), ], ) def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_declared_level( - gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue] + gateway: Gateway, kind: str, target: str, fields: Mapping[str, JsonValue], expected: Expected ) -> None: marker: Final = new_marker() response_format: Final[dict[str, JsonValue]] = { @@ -1844,10 +2040,116 @@ def test_a_deployment_opted_into_bedrock_native_output_keeps_it_beside_its_decla assert response.status_code == 200, response.text[:2000] assert answer(marker) in response.text, response.text[:2000] assert [(request.method, request.target) for request in received] == [("POST", target)] - body: Final = body_of(received[0]) - assert _bedrock_native_schema(body, invoke=target == _INVOKE) == _NATIVE_OUTPUT_SCHEMA, body - assert {"toolConfig", "tools"}.isdisjoint(body), body - assert sent_effort(received[0]) == "high", body + assert body_of(received[0]) == expected(marker) + + +_NATIVE_OUTPUT_REJECTION: Final = ( + "The model returned the following errors: output_config.format: Extra inputs are not permitted" +) + + +def _rejecting_native_output(request: Request) -> Reply: + native: Final = _bedrock_native_schema(body_of(request), invoke=path_of(request).endswith("/invoke")) + if native is None: + return provider_reply(request) + return Reply(status=400, body=json.dumps({"message": _NATIVE_OUTPUT_REJECTION}).encode()) + + +@pytest.mark.parametrize( + ("kind", "target", "rejected", "retried"), + [ + pytest.param( + "bedrock_converse", _CONVERSE, _CONVERSE_WITH_NATIVE_OUTPUT, _CONVERSE_WITH_JSON_TOOL, id="converse" + ), + pytest.param("bedrock_invoke", _INVOKE, _INVOKE_WITH_NATIVE_OUTPUT, _INVOKE_WITH_JSON_TOOL, id="invoke"), + ], +) +def test_the_retry_after_bedrock_rejects_native_output_keeps_the_declared_level( + gateway: Gateway, kind: str, target: str, rejected: Expected, retried: Expected +) -> None: + marker: Final = new_marker() + response_format: Final[dict[str, JsonValue]] = { + "type": "json_schema", + "json_schema": {"name": "answer", "schema": _NATIVE_OUTPUT_SCHEMA}, + } + with wire_server(_rejecting_native_output) as wire, gateway.scenario() as scenario: + model: Final = deployment(scenario, kind, wire, _CLAUDE_INFO_WITH_NATIVE_OUTPUT) + response: Final = gateway.request( + "POST", CHAT, chat_request(model, marker, response_format=response_format, thinking=_BUDGET) + ) + received: Final = wire.drain() + assert response.status_code == 200, response.text[:2000] + assert [(request.method, request.target) for request in received] == [("POST", target), ("POST", target)] + assert [body_of(request) for request in received] == [rejected(marker), retried(marker)] + + +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_TOKEN: Final = "/_oauth/token" +_VERTEX_RAW_PREDICT: Final = ( + f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}/publishers/anthropic/models/{CLAUDE}:rawPredict" +) +_VERTEX_ACCESS_TOKEN: Final = json.dumps({"access_token": "scripted", "token_type": "Bearer", "expires_in": 3600}) + + +def _vertex(expected: Expected) -> Expected: + return _with_fields(_without(expected, "model"), {"anthropic_version": "vertex-2023-10-16"}) + + +def _vertex_reply(request: Request) -> Reply: + if request.target == _VERTEX_TOKEN: + return Reply(body=_VERTEX_ACCESS_TOKEN.encode()) + marker: Final = marker_of(request) + return Reply(body=message_json(f"msg_{marker}", CLAUDE, answer(marker))) + + +@pytest.mark.parametrize( + ("route", "fields", "expected", "status"), + [ + pytest.param( + CHAT, + {"reasoning_effort": "medium"}, + _vertex(_anthropic_chat("high", {"type": "adaptive", "display": "summarized"})), + 200, + id="chat-medium-to-high", + ), + pytest.param( + MESSAGES, + {"thinking": _BUDGET}, + _vertex(_native_messages({"effort": "high"}, _ADAPTIVE)), + 200, + id="messages-budget-to-high", + ), + pytest.param(CHAT, {"thinking": _FALSE_BUDGET}, None, 400, id="chat-false-budget-400"), + ], +) +def test_a_vertex_ai_claude_deployment_gets_one_of_its_declared_levels( + gateway: Gateway, route: str, fields: dict[str, JsonValue], expected: Expected | None, status: int +) -> None: + marker: Final = new_marker() + with wire_server(_vertex_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"vertex_ai/{CLAUDE}", + api_base=wire.url, + api_key=None, + vertex_project=_VERTEX_PROJECT, + vertex_location=_VERTEX_LOCATION, + vertex_credentials=service_account_json(_VERTEX_PROJECT, wire.url), + model_info=dict(_CLAUDE_INFO), + ) + body: Final = ( + chat_request(model, marker, **fields) if route == CHAT else messages_request(model, marker, **fields) + ) + response: Final = gateway.request("POST", route, body) + received: Final = tuple(request for request in wire.drain() if request.target != _VERTEX_TOKEN) + assert response.status_code == status, response.text[:2000] + if expected is None: + assert _NON_NUMERIC_BUDGET in response.text, response.text[:2000] + assert received == (), [request.body[:300] for request in received] + return + assert answer(marker) in response.text, response.text[:2000] + assert [(request.method, request.target) for request in received] == [("POST", _VERTEX_RAW_PREDICT)] + assert body_of(received[0]) == expected(marker) def _single(received: tuple[Request, ...], marker: str) -> Request: diff --git a/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py b/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py new file mode 100644 index 00000000000..d0235310aaa --- /dev/null +++ b/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py @@ -0,0 +1,67 @@ +"""An SDK caller that reuses one ``metadata=`` dict across ``litellm.completion`` calls to a Bedrock deployment whose +``guardrailConfig`` has ``trace: enabled`` is billed, on each call, for that call's guardrail trace only, and the +caller's dict is left as the caller wrote it (LIT-10239)""" + +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from integration._support.client import eventually +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm + +_MODEL: Final = "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0" +_PROMPT_TOKENS: Final = 5 +_COMPLETION_TOKENS: Final = 1 +_USAGE: Final[Mapping[str, int]] = {"contentPolicyUnits": 1, "topicPolicyUnits": 1} +_GUARDRAIL_IDS: Final = ("guardrail-one", "guardrail-two") + + +def _bedrock(request: Request) -> Reply: + """Bedrock's Converse runtime: a one-token answer and the trace of the guardrail the request named.""" + guardrail_id: Final = json.loads(request.body)["guardrailConfig"]["guardrailIdentifier"] + body: Final[dict[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6}, + "metrics": {"latencyMs": 3}, + "trace": {"guardrail": {"inputAssessment": {guardrail_id: {"invocationMetrics": {"usage": dict(_USAGE)}}}}}, + } + return Reply(body=json.dumps(body).encode()) + + +def _guardrail_cost() -> float: + rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"] + return sum(units * rates[counter] for counter, units in _USAGE.items()) + + +def _llm_cost() -> float: + prices: Final = litellm.model_cost[_MODEL.removeprefix("bedrock/")] + return _PROMPT_TOKENS * prices["input_cost_per_token"] + _COMPLETION_TOKENS * prices["output_cost_per_token"] + + +def test_each_call_sharing_one_metadata_dict_is_billed_its_own_trace_only() -> None: + billed: list[float] = [] # mutable-ok: the success callback appends what each call was billed + litellm.success_callback.append( + lambda kwargs, response, start, end: billed.append(kwargs["standard_logging_object"]["response_cost"]) + ) + shared_metadata: Final = {"tag": "x"} + with wire_server(_bedrock) as wire: + for guardrail_id in _GUARDRAIL_IDS: + litellm.completion( + model=_MODEL, + messages=[{"role": "user", "content": "hello"}], + api_base=wire.url, + aws_access_key_id="AKIASCRIPTED", + aws_secret_access_key="scripted", + aws_region_name="us-east-1", + guardrailConfig={"guardrailIdentifier": guardrail_id, "guardrailVersion": "1", "trace": "enabled"}, + metadata=shared_metadata, + ) + costs: Final = eventually(lambda: tuple(billed), lambda seen: len(seen) == len(_GUARDRAIL_IDS), seconds=30) + + assert costs == pytest.approx((_llm_cost() + _guardrail_cost(),) * len(_GUARDRAIL_IDS)) + assert shared_metadata == {"tag": "x"} diff --git a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py index 747adeee41d..e7508a2670e 100644 --- a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py +++ b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py @@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm import Router +from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments from litellm.types.router import DeploymentTypedDict _MODEL: Final = "gpt-5.4" @@ -118,3 +119,30 @@ async def test_cooldown_skips_only_the_failed_duplicate_model_deployment() -> No "/primary/chat/completions", "/backup/chat/completions", ) + + +@pytest.mark.parametrize( + "router_kwargs", + [{"fallbacks": [{"primary": "backup"}]}, {"default_fallbacks": "backup"}], + ids=["string_fallback_target", "string_default_fallbacks"], +) +@pytest.mark.asyncio +async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(router_kwargs: dict[str, object]) -> None: + with wire_server(_peer) as wire: + router: Final = Router( + model_list=[ + _deployment("primary", f"{wire.url}/primary", "primary-deployment"), + _deployment("backup", f"{wire.url}/backup", "backup-deployment"), + ], + allowed_fails=0, + num_retries=0, + **router_kwargs, + ) + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="primary", + messages=[{"role": "user", "content": f"cooldown {uuid.uuid4().hex}"}], + ) + cooled: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + + assert tuple(cooled) == ("primary-deployment",) diff --git a/tests/integration/security/test_hide_secrets_tool_traffic.py b/tests/integration/security/test_hide_secrets_tool_traffic.py index 42fdc0f5a44..60b49520b1f 100644 --- a/tests/integration/security/test_hide_secrets_tool_traffic.py +++ b/tests/integration/security/test_hide_secrets_tool_traffic.py @@ -413,3 +413,71 @@ def test_resent_tool_history_is_scanned_once(single_worker: Rig, shape: str) -> second: Final = timed(lambda: single_worker.proxy.request("POST", path, following)) assert first.value.status_code == 200 and second.value.status_code == 200, (first.value.text, second.value.text) assert_within_factor_of_control(second.seconds, control=first.seconds, factor=0.5) + + +@pytest.mark.parametrize("placement", ("guardrail", "global")) +def test_skip_tool_message_in_guardrail_forwards_tool_results(tmp_path: Path, placement: str) -> None: + def configure(configuration: dict[str, JsonValue]) -> None: + guardrails: Final = configuration["guardrails"] + settings: Final = configuration["litellm_settings"] + assert isinstance(guardrails, list) and isinstance(guardrails[0], dict) and isinstance(settings, dict) + params: Final = guardrails[0]["litellm_params"] + assert isinstance(params, dict) + (params if placement == "guardrail" else settings)["skip_tool_message_in_guardrail"] = True + + with tool_rig(tmp_path, configure=configure) as rig: + rig.wire.drain() + replies: Final = ( + rig.proxy.request( + "POST", + "/v1/messages", + { + "model": "claude", + "max_tokens": 32, + "cache": {"no-cache": True}, + "messages": [ + {"role": "user", "content": TEXT}, + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "tool_1", "name": "read_file", "input": {"env": TEXT}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "tool_1", "content": TEXT}], + }, + ], + }, + ), + rig.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": "gpt", + "cache": {"no-cache": True}, + "messages": [ + {"role": "user", "content": TEXT}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "read_file", "arguments": json.dumps({"env": TEXT})}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": TEXT}, + ], + }, + ), + ) + assert all(reply.status_code == 200 for reply in replies), [reply.text for reply in replies] + sent: Final = rig.wire.drain() + assert len(sent) == 2 + for request in sent: + messages = json.loads(request.body)["messages"] + assert SECRET not in json.dumps(messages[:2]) + assert SECRET in json.dumps(messages[2]) diff --git a/tests/integration/security/test_managed_file_listing.py b/tests/integration/security/test_managed_file_listing.py new file mode 100644 index 00000000000..8c736b2b2bf --- /dev/null +++ b/tests/integration/security/test_managed_file_listing.py @@ -0,0 +1,263 @@ +import json +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from itertools import count, islice +from pathlib import Path +from types import MappingProxyType +from typing import Final +from urllib.parse import parse_qs, urlsplit + +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from anthropic import Anthropic +from openai import OpenAI +from pydantic import BaseModel, ConfigDict, JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +UNMANAGED_FILE: Final = "file-uploaded-straight-to-the-provider" +UPLOAD: Final = b"isolation check\n" + + +@dataclass(frozen=True, slots=True) +class Deployment: + provider: str + model: str + litellm_model: str + api_key: str + api_path: str + + +OPENAI: Final = Deployment("openai", "gpt-4.1-managed-listing", "openai/gpt-4.1", "sk-managed-listing", "/v1") +ANTHROPIC: Final = Deployment( + "anthropic", "claude-managed-listing", "anthropic/claude-haiku-5-5", "sk-ant-managed-listing", "" +) +DEPLOYMENTS: Final = (pytest.param(OPENAI, id="openai"), pytest.param(ANTHROPIC, id="anthropic")) + + +@dataclass(frozen=True, slots=True) +class ListCall: + path: str + params: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + headers: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + + +LIST_FORMS: Final = ( + pytest.param(lambda d: ListCall(f"/{d.provider}/v1/files"), id="provider-path"), + pytest.param(lambda d: ListCall("/v1/files", params={"provider": d.provider}), id="provider-query"), + pytest.param(lambda d: ListCall("/v1/files", headers={"custom-llm-provider": d.provider}), id="provider-header"), + pytest.param(lambda d: ListCall("/v1/files", params={"target_model_names": d.model}), id="target-model-names"), + pytest.param(lambda d: ListCall("/v1/files", headers={"x-litellm-model": d.model}), id="model-header"), +) + + +class FilePage(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + data: list[dict[str, JsonValue]] + has_more: bool + + @property + def ids(self) -> list[str]: + return [str(entry["id"]) for entry in self.data] + + +def _openai_file(file_id: str) -> dict[str, JsonValue]: + return { + "id": file_id, + "object": "file", + "bytes": len(UPLOAD), + "created_at": 1700000000, + "filename": "notes.txt", + "purpose": "user_data", + "status": "processed", + } + + +def _anthropic_file(file_id: str) -> dict[str, JsonValue]: + return { + "id": file_id, + "type": "file", + "filename": "notes.txt", + "mime_type": "text/plain", + "size_bytes": len(UPLOAD), + "created_at": "2026-01-01T00:00:00Z", + "downloadable": False, + } + + +def _reply(body: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps(body).encode()) + + +def _list_reply(provider: str, file_ids: tuple[str, ...]) -> Reply: + if provider == ANTHROPIC.provider: + return _reply( + { + "data": [_anthropic_file(file_id) for file_id in file_ids], + "has_more": False, + "first_id": file_ids[0], + "last_id": file_ids[-1], + } + ) + return _reply({"object": "list", "data": [_openai_file(file_id) for file_id in file_ids]}) + + +def _responder(provider_files: Mapping[str, tuple[str, ...]]) -> Callable[[Request], Reply]: + uploads: Final = {provider: count() for provider in provider_files} + + def respond(request: Request) -> Reply: + target: Final = urlsplit(request.target) + path: Final = target.path + provider: Final = ANTHROPIC.provider if "x-api-key" in request.headers else OPENAI.provider + files: Final = provider_files[provider] + if (request.method, path) == ("GET", "/v1/models"): + return _reply({"object": "list", "data": []}) + if (request.method, path) == ("POST", "/v1/files"): + uploaded: Final = files[next(uploads[provider])] + return _reply(_anthropic_file(uploaded) if provider == ANTHROPIC.provider else _openai_file(uploaded)) + if (request.method, path) == ("GET", "/v1/files"): + limit: Final = int(parse_qs(target.query).get("limit", ["10000"])[0]) + return _list_reply(provider, (UNMANAGED_FILE, *reversed(files))[:limit]) + return Reply(status=404, body=json.dumps({"error": {"message": f"unscripted {request.target}"}}).encode()) + + return respond + + +def _config(wire: Wire) -> dict[str, JsonValue]: + return { + "model_list": [ + { + "model_name": deployment.model, + "litellm_params": { + "model": deployment.litellm_model, + "api_key": deployment.api_key, + "api_base": f"{wire.url}{deployment.api_path}", + }, + } + for deployment in (OPENAI, ANTHROPIC) + ], + "litellm_settings": {"require_managed_files": True}, + "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL"}, + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + owner: str + intruder: str + files_by_provider: Mapping[str, tuple[str, ...]] + every_file: tuple[str, ...] + + @property + def base_url(self) -> str: + return str(self.proxy.client.base_url).rstrip("/") + + def list(self, call: ListCall, key: str | None = None) -> FilePage: + response: Final = self.proxy.request("GET", call.path, key=key, params=call.params, headers=call.headers) + assert response.status_code == 200, response.text + return FilePage.model_validate_json(response.text) + + +def _upload(proxy: Gateway, key: str, deployments: tuple[Deployment, ...]) -> str: + with OpenAI(base_url=f"{str(proxy.client.base_url).rstrip('/')}/openai/v1", api_key=key, max_retries=0) as client: + uploaded: Final = client.files.with_raw_response.create( + file=("notes.txt", UPLOAD, "text/plain"), + purpose="user_data", + extra_body={"target_model_names": ",".join(deployment.model for deployment in deployments)}, + ) + assert uploaded.status_code == 200, uploaded.text + return uploaded.parse().id + + +@contextmanager +def _rig(directory: Path) -> Iterator[Rig]: + build: Final = uuid.uuid4().hex + provider_files: Final = { + OPENAI.provider: (f"file-older-{build}", f"file-newer-{build}"), + ANTHROPIC.provider: (f"file_older_{build}", f"file_newer_{build}"), + } + with gateway_from_environment() as shared, wire_server(_responder(provider_files)) as wire: + config: Final = directory / "config.yaml" + config.write_text(yaml.safe_dump(_config(wire))) + with owned_proxy(shared, directory, {}, config=config) as proxy, proxy.scenario() as scenario: + owner: Final = scenario.key(user_id=scenario.user()) + intruder: Final = scenario.key(user_id=scenario.user()) + on_both: Final = _upload(proxy, owner, (OPENAI, ANTHROPIC)) + on_openai: Final = _upload(proxy, owner, (OPENAI,)) + on_anthropic: Final = _upload(proxy, owner, (ANTHROPIC,)) + yield Rig( + proxy, + owner, + intruder, + MappingProxyType( + {OPENAI.provider: (on_openai, on_both), ANTHROPIC.provider: (on_anthropic, on_both)} + ), + (on_anthropic, on_openai, on_both), + ) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _rig(tmp_path_factory.mktemp("managed-file-listing")) as built: + yield built + + +@pytest.mark.parametrize("deployment", DEPLOYMENTS) +@pytest.mark.parametrize("form", LIST_FORMS) +def test_another_key_lists_none_of_the_files_on_the_shared_key( + rig: Rig, deployment: Deployment, form: Callable[[Deployment], ListCall] +) -> None: + page: Final = rig.list(form(deployment), key=rig.intruder) + assert (page.ids, page.has_more) == ([], False) + + +@pytest.mark.parametrize("deployment", DEPLOYMENTS) +@pytest.mark.parametrize("form", LIST_FORMS) +def test_owner_lists_their_own_managed_files_on_that_provider_or_model_newest_first( + rig: Rig, deployment: Deployment, form: Callable[[Deployment], ListCall] +) -> None: + page: Final = rig.list(form(deployment), key=rig.owner) + assert (page.ids, page.has_more) == (list(rig.files_by_provider[deployment.provider]), False) + + +def test_owner_lists_every_managed_file_when_no_provider_or_model_is_named(rig: Rig) -> None: + page: Final = rig.list(ListCall("/v1/files"), key=rig.owner) + assert (page.ids, page.has_more) == (list(rig.every_file), False) + + +@pytest.mark.parametrize("cursor_param", ["after", "after_id"]) +def test_owner_pages_through_a_provider_list_with_limit_and_a_cursor(rig: Rig, cursor_param: str) -> None: + first: Final = rig.list(ListCall("/openai/v1/files", params={"limit": "1"}), key=rig.owner) + second: Final = rig.list( + ListCall("/openai/v1/files", params={"limit": "1", cursor_param: first.ids[0]}), key=rig.owner + ) + expected: Final = rig.files_by_provider[OPENAI.provider] + assert [(page.ids, page.has_more) for page in (first, second)] == [ + ([expected[0]], True), + ([expected[1]], False), + ] + + +def test_anthropic_sdk_auto_pages_the_owners_files_to_the_end(rig: Rig) -> None: + with Anthropic(base_url=f"{rig.base_url}/anthropic", api_key=rig.owner, max_retries=0) as client: + listed: Final = [file.id for file in islice(client.beta.files.list(limit=1), 5)] + assert listed == list(rig.files_by_provider[ANTHROPIC.provider]) + + +def test_anthropic_sdk_pages_backwards_to_the_newest_file_and_stops(rig: Rig) -> None: + newest_first: Final = rig.files_by_provider[ANTHROPIC.provider] + with Anthropic(base_url=f"{rig.base_url}/anthropic", api_key=rig.owner, max_retries=0) as client: + listed: Final = [file.id for file in islice(client.beta.files.list(limit=1, before_id=newest_first[-1]), 5)] + assert listed == list(reversed(newest_first[:-1])) + + +def test_master_key_still_lists_the_files_on_the_shared_provider_key(rig: Rig) -> None: + listed: Final = rig.list(ListCall(f"/{ANTHROPIC.provider}/v1/files")).ids + assert UNMANAGED_FILE in listed diff --git a/tests/integration/spend/fixtures/daily_activity_team.json b/tests/integration/spend/fixtures/daily_activity_team.json index 62377c07432..8f4956dd7fe 100644 --- a/tests/integration/spend/fixtures/daily_activity_team.json +++ b/tests/integration/spend/fixtures/daily_activity_team.json @@ -1 +1 @@ -{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"team-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{"team_alias":"Usage Team"},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"team-1":5}}} +{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"team-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{"team_alias":"Usage Team"},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}},"guardrails":{}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"team-1":5}}} diff --git a/tests/integration/spend/fixtures/daily_activity_user.json b/tests/integration/spend/fixtures/daily_activity_user.json index c4302ee7045..ae05e661071 100644 --- a/tests/integration/spend/fixtures/daily_activity_user.json +++ b/tests/integration/spend/fixtures/daily_activity_user.json @@ -1 +1 @@ -{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"user-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"user-1":5}}} +{"results":[{"date":"2026-06-01","metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"breakdown":{"mcp_servers":{},"models":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"model_groups":{"model-ptu":{"metrics":{"spend":1000.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":0,"completion_tokens":0,"cache_read_input_tokens":0,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":0,"successful_requests":0,"failed_requests":0,"api_requests":0,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{}},"model-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}},"model-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}},"model-popular":{"metrics":{"spend":270.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":9,"completion_tokens":6,"cache_read_input_tokens":3,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":15,"successful_requests":3,"failed_requests":0,"api_requests":3,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"providers":{"provider-a":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"endpoints":{"/v1/chat/completions":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}}}}},"api_keys":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}},"entities":{"user-1":{"metrics":{"spend":1273.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1014,"completion_tokens":10,"cache_read_input_tokens":6,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1024,"successful_requests":5,"failed_requests":0,"api_requests":5,"total_response_time_ms":0,"timed_requests":0},"metadata":{},"api_key_breakdown":{"key-a":{"metrics":{"spend":100.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":2,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":4,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-a","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-b":{"metrics":{"spend":90.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":3,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":5,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-b","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-c":{"metrics":{"spend":80.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":4,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":6,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-c","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-cache":{"metrics":{"spend":2.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":1000,"completion_tokens":2,"cache_read_input_tokens":1,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":1002,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"alias-cache","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":true}},"key-target":{"metrics":{"spend":1.0,"flat_cost":0.0,"ptu_hours":0.0,"prompt_tokens":5,"completion_tokens":2,"cache_read_input_tokens":2,"cache_creation_input_tokens":0,"compression_saved_tokens":0,"compression_savings_spend":0.0,"prompt_caching_savings_spend":0.0,"gateway_injected_caching_savings_spend":0.0,"autorouter_savings_spend":0.0,"total_tokens":7,"successful_requests":1,"failed_requests":0,"api_requests":1,"total_response_time_ms":0,"timed_requests":0},"metadata":{"key_alias":"deleted-target","team_id":"team-1","user_id":"user-1","user_email":"user@example.com","key_exists":false}}}}},"guardrails":{}}}],"metadata":{"total_spend":1273.0,"total_flat_cost":0.0,"total_ptu_hours":0.0,"total_prompt_tokens":1014,"total_completion_tokens":10,"total_tokens":1024,"total_api_requests":5,"total_successful_requests":5,"total_failed_requests":0,"total_cache_read_input_tokens":6,"total_cache_creation_input_tokens":0,"total_compression_saved_tokens":0,"total_compression_savings_spend":0.0,"total_prompt_caching_savings_spend":0.0,"total_gateway_injected_caching_savings_spend":0.0,"total_autorouter_savings_spend":0.0,"total_response_time_ms":0,"total_timed_requests":0,"page":1,"total_pages":1,"has_more":false,"api_key_limit":100,"total_api_keys":5,"entity_total_api_keys":{"user-1":5}}} diff --git a/tests/integration/spend/golden/daily_activity_team_aggregated.json b/tests/integration/spend/golden/daily_activity_team_aggregated.json index 2956adda17a..91a5d33391b 100644 --- a/tests/integration/spend/golden/daily_activity_team_aggregated.json +++ b/tests/integration/spend/golden/daily_activity_team_aggregated.json @@ -26,7 +26,7 @@ }, "results": [ { - "breakdown": { + "breakdown": {"guardrails": {}, "api_keys": { "key-a": { "metadata": { diff --git a/tests/integration/spend/golden/daily_activity_team_paginated.json b/tests/integration/spend/golden/daily_activity_team_paginated.json index ca7239f8fa9..4d91b0da209 100644 --- a/tests/integration/spend/golden/daily_activity_team_paginated.json +++ b/tests/integration/spend/golden/daily_activity_team_paginated.json @@ -26,7 +26,7 @@ }, "results": [ { - "breakdown": { + "breakdown": {"guardrails": {}, "api_keys": { "key-a": { "metadata": { diff --git a/tests/integration/spend/golden/daily_activity_user_aggregated.json b/tests/integration/spend/golden/daily_activity_user_aggregated.json index 261ac22b2cb..185d333965a 100644 --- a/tests/integration/spend/golden/daily_activity_user_aggregated.json +++ b/tests/integration/spend/golden/daily_activity_user_aggregated.json @@ -26,7 +26,7 @@ }, "results": [ { - "breakdown": { + "breakdown": {"guardrails": {}, "api_keys": { "key-a": { "metadata": { diff --git a/tests/integration/spend/golden/daily_activity_user_paginated.json b/tests/integration/spend/golden/daily_activity_user_paginated.json index f1fb7ba50a4..00389289ebe 100644 --- a/tests/integration/spend/golden/daily_activity_user_paginated.json +++ b/tests/integration/spend/golden/daily_activity_user_paginated.json @@ -26,7 +26,7 @@ }, "results": [ { - "breakdown": { + "breakdown": {"guardrails": {}, "api_keys": { "key-a": { "metadata": { diff --git a/tests/integration/spend/test_background_response_billed_once.py b/tests/integration/spend/test_background_response_billed_once.py index b207e46204a..5f948fac3ab 100644 --- a/tests/integration/spend/test_background_response_billed_once.py +++ b/tests/integration/spend/test_background_response_billed_once.py @@ -1,6 +1,7 @@ import asyncio import base64 import json +import math import uuid from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor @@ -10,9 +11,18 @@ from typing import Final import httpx import openai import pytest -from integration._support.client import Gateway, Scenario, eventually, held, list_value, object_value, string_value +from integration._support.client import ( + Gateway, + Scenario, + delete_key_if_present, + eventually, + held, + list_value, + object_value, + string_value, +) from integration._support.database import read_rows, write_rows -from integration._support.otlp_sink import decode_spans +from integration._support.otlp_sink import Span, decode_spans from integration._support.upstream import delete_scenario, register_scenario from integration._support.wire import Request, wire_server from integration.caching.response_cache_case import entry_key, redis_store @@ -859,7 +869,7 @@ LANGFUSE_COST: Final = "langfuse.observation.cost_details" _LANGFUSE_COST_DETAILS: Final = TypeAdapter(dict[str, float]) -def _langfuse_callback(host: str) -> dict[str, JsonValue]: +def _langfuse_callback(host: str, **callback_vars: JsonValue) -> dict[str, JsonValue]: return { "callback_name": "langfuse", "callback_type": "success", @@ -867,12 +877,17 @@ def _langfuse_callback(host: str) -> dict[str, JsonValue]: "langfuse_public_key": LANGFUSE_PUBLIC_KEY, "langfuse_secret_key": LANGFUSE_SECRET_KEY, "langfuse_host": host, + **callback_vars, }, } +def _logging_to(langfuse_url: str, **callback_vars: JsonValue) -> dict[str, JsonValue]: + return {"logging": [_langfuse_callback(langfuse_url, **callback_vars)]} + + def _key_with_its_own_logging(scenario: Scenario, langfuse_url: str) -> str: - return scenario.key(metadata={"logging": [_langfuse_callback(langfuse_url)]}) + return scenario.key(metadata=_logging_to(langfuse_url)) def _key_in_a_team_with_logging(scenario: Scenario, langfuse_url: str) -> str: @@ -892,17 +907,33 @@ def _trace_exports(langfuse: Recorder) -> tuple[Request, ...]: return tuple(request for request in langfuse.requests() if request.target.startswith(LANGFUSE_TRACES_PATH)) -def _exported_costs(langfuse: Recorder) -> tuple[float, ...]: - spans: Final = tuple( +def _exported_spans(langfuse: Recorder) -> tuple[Span, ...]: + return tuple( chain.from_iterable( decode_spans(export.body, export.headers.get("content-type", "")) for export in _trace_exports(langfuse) ) ) - return tuple( - _LANGFUSE_COST_DETAILS.validate_json(string_value(span["attributes"][LANGFUSE_COST]))["total"] - for span in spans - if LANGFUSE_COST in span["attributes"] - ) + + +def _span_cost(span: Span) -> float | None: + cost_details: Final = span["attributes"].get(LANGFUSE_COST) + return None if cost_details is None else _LANGFUSE_COST_DETAILS.validate_json(string_value(cost_details))["total"] + + +def _exported_costs(langfuse: Recorder) -> tuple[float, ...]: + return tuple(cost for cost in map(_span_cost, _exported_spans(langfuse)) if cost is not None) + + +def _priced_spans(langfuse: Recorder) -> tuple[Span, ...]: + return tuple(span for span in _exported_spans(langfuse) if (_span_cost(span) or 0.0) > 0) + + +def _billed_once(langfuse: Recorder, seconds: float) -> Span: + priced: Final = eventually(lambda: _priced_spans(langfuse), lambda spans: len(spans) > 0, seconds=seconds) + costs: Final = tuple(_span_cost(span) or 0.0 for span in priced) + assert len(costs) == 1, costs + assert math.isclose(costs[0], PRICE), costs + return priced[0] @pytest.mark.timeout(240) @@ -925,3 +956,133 @@ def test_the_billed_cost_reaches_the_langfuse_project_set_on_the_creating_key_or assert [cost for cost in exported if cost > 0] == [pytest.approx(PRICE)] credentials: Final = base64.b64encode(f"{LANGFUSE_PUBLIC_KEY}:{LANGFUSE_SECRET_KEY}".encode()).decode() assert {export.headers.get("authorization") for export in _trace_exports(langfuse)} == {f"Basic {credentials}"} + + +@pytest.mark.timeout(240) +def test_the_billed_cost_reaches_the_creating_keys_langfuse_project_and_never_the_readers(gateway: Gateway) -> None: + with ( + wire_server(langfuse_sink) as creator_wire, + wire_server(langfuse_sink) as reader_wire, + gateway.scenario() as scenario, + ): + creator_langfuse: Final = Recorder(creator_wire) + reader_langfuse: Final = Recorder(reader_wire) + team_id: Final = scenario.team() + user_id: Final = scenario.member(team_id) + creator: Final = scenario.key(team_id=team_id, user_id=user_id, metadata=_logging_to(creator_wire.url)) + reader: Final = scenario.key(team_id=team_id, user_id=user_id, metadata=_logging_to(reader_wire.url)) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, creator, upstream.model) + reads: Final = tuple(read_response(gateway, reader, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + _billed_once(creator_langfuse, SPEND_LANDING_SECONDS) + eventually(lambda: _trace_exports(reader_langfuse), lambda exports: len(exports) > 0) + assert _priced_spans(reader_langfuse) == () + + +MESSAGE_LOGGING: Final[Mapping[str, Mapping[str, JsonValue]]] = { + "logged": {}, + "redacted": {"turn_off_message_logging": "true"}, +} + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("message_logging", MESSAGE_LOGGING.keys()) +def test_the_billed_observation_carries_the_response_text_only_when_the_creating_key_keeps_message_logging( + gateway: Gateway, message_logging: str +) -> None: + with wire_server(langfuse_sink) as wire, gateway.scenario() as scenario: + langfuse: Final = Recorder(wire) + key: Final = scenario.key(metadata=_logging_to(wire.url, **MESSAGE_LOGGING[message_logging])) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, key, upstream.model) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + billed: Final = _billed_once(langfuse, SPEND_LANDING_SECONDS) + assert ("pong" in json.dumps(billed["attributes"])) is (message_logging == "logged"), billed["attributes"] + + +def _generated_key(scenario: Scenario, **fields: JsonValue) -> str: + key: Final = string_value(scenario.gateway.post("/key/generate", fields)["key"]) + scenario.cleanups.callback(delete_key_if_present, scenario.gateway, key) + return key + + +def _delete_team_if_present(gateway: Gateway, team_id: str) -> None: + if gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)): + gateway.post("/team/delete", {"team_ids": [team_id]}) + + +Retirement = Callable[[Scenario, str, str], tuple[str, Callable[[], object]]] + + +def _rotated_with_grace_then_relogged( + scenario: Scenario, billed_url: str, stale_url: str +) -> tuple[str, Callable[[], object]]: + key: Final = _generated_key(scenario, metadata=_logging_to(stale_url)) + + def rotate() -> None: + rotated: Final = string_value( + scenario.gateway.post("/key/regenerate", {"key": key, "grace_period": "1h"})["key"] + ) + scenario.cleanups.callback(delete_key_if_present, scenario.gateway, rotated) + scenario.gateway.post("/key/update", {"key": rotated, "metadata": _logging_to(billed_url)}) + + return key, rotate + + +def _deleted_key(scenario: Scenario, billed_url: str, _stale_url: str) -> tuple[str, Callable[[], object]]: + key: Final = _generated_key(scenario, metadata=_logging_to(billed_url)) + return key, lambda: scenario.gateway.post("/key/delete", {"keys": [key]}) + + +def _deleted_team(scenario: Scenario, billed_url: str, _stale_url: str) -> tuple[str, Callable[[], object]]: + team_id: Final = string_value( + scenario.gateway.post("/team/new", {"team_alias": f"integration-{uuid.uuid4().hex}"})["team_id"] + ) + scenario.cleanups.callback(_delete_team_if_present, scenario.gateway, team_id) + added: Final = scenario.gateway.request("POST", f"/team/{team_id}/callback", _langfuse_callback(billed_url)) + assert added.status_code == 200, added.text + key: Final = _generated_key(scenario, team_id=team_id) + return key, lambda: scenario.gateway.post("/team/delete", {"team_ids": [team_id]}) + + +RETIREMENTS: Final[Mapping[str, Retirement]] = { + "key_rotated_with_grace": _rotated_with_grace_then_relogged, + "key_deleted": _deleted_key, + "team_deleted": _deleted_team, +} + + +@pytest.mark.timeout(240) +@pytest.mark.parametrize("retired", RETIREMENTS.keys()) +def test_a_retired_creators_billed_cost_reaches_its_own_langfuse_project(gateway: Gateway, retired: str) -> None: + with ( + wire_server(langfuse_sink) as billed_wire, + wire_server(langfuse_sink) as stale_wire, + gateway.scenario() as scenario, + ): + billed_langfuse: Final = Recorder(billed_wire) + stale_langfuse: Final = Recorder(stale_wire) + key, retire = RETIREMENTS[retired](scenario, billed_wire.url, stale_wire.url) + upstream: Final = scripted_upstream(scenario, background_routes("queued")) + response_id: Final = create_background_response(scenario, key, upstream.model) + retire() + upstream.script(background_routes("completed")) + _billed_once(billed_langfuse, ONE_POLL_SECONDS + SPEND_LANDING_SECONDS) + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] + assert _priced_spans(stale_langfuse) == () + + +@pytest.mark.timeout(240) +def test_a_creator_whose_langfuse_keys_resolve_to_nothing_is_still_billed_once(gateway: Gateway) -> None: + unset: Final = f"os.environ/INTEGRATION_UNSET_LANGFUSE_{uuid.uuid4().hex.upper()}" + with wire_server(langfuse_sink) as wire, gateway.scenario() as scenario: + key: Final = scenario.key(metadata=_logging_to(wire.url, langfuse_public_key=unset, langfuse_secret_key=unset)) + upstream: Final = scripted_upstream(scenario, background_routes("completed")) + response_id: Final = create_background_response(scenario, key, upstream.model) + reads: Final = tuple(read_response(gateway, key, response_id) for _ in range(READS)) + assert reads in (settling_reads("completed", READS), free_reads("completed", READS)) + assert_charged_exactly(lambda: key_spend(key), PRICE) + assert math.isclose(number(one_charged_row(upstream.model_id)["spend"]), PRICE) + assert [row["status"] for row in managed_rows(response_id)] == ["completed"] diff --git a/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py b/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py index 33baf9d0e2e..1c315a6e694 100644 --- a/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py +++ b/tests/integration/spend/test_daily_activity_aggregated_breakdowns.py @@ -243,3 +243,183 @@ async def test_get_daily_activity_aggregated_model_group_rollups_fall_back_to_mo assert object_value(object_value(models["gpt-5"])["metrics"])["spend"] == 10.0 finally: _clean(day) + + +def _seed_guardrail_units(day: str, rows: Sequence[tuple[str, str, str, float | None, int, int]]) -> None: + for guardrail_id, api_key, usage_unit, cost, units, untracked_units in rows: + write_rows( + 'INSERT INTO "LiteLLM_DailyGuardrailUsageUnits" (guardrail_id, date, team_id, api_key, usage_unit,' + " units, cost, untracked_units, updated_at) VALUES (%s, %s, '', %s, %s, %s, %s, %s, now())", + ( + guardrail_id, + day, + api_key, + usage_unit, + str(units), + None if cost is None else str(cost), + str(untracked_units), + ), + ) + + +def _guardrail_spend(breakdown: dict[str, JsonValue]) -> dict[str, tuple[float, dict[str, float]]]: + return { + guardrail_id: ( + float(object_value(object_value(entry)["metrics"])["spend"]), + { + api_key: float(object_value(object_value(key_entry)["metrics"])["spend"]) + for api_key, key_entry in object_value(object_value(entry)["api_key_breakdown"]).items() + }, + ) + for guardrail_id, entry in object_value(breakdown["guardrails"]).items() + } + + +def _day_breakdown(gateway: Gateway, day: str, **params: str) -> dict[str, JsonValue]: + results: Final = _RESULTS.validate_python(_activity(gateway, day, **params)["results"]) + assert len(results) == 1 + return object_value(object_value(results[0])["breakdown"]) + + +@pytest.mark.asyncio +async def test_get_daily_activity_aggregated_breaks_down_guardrail_spend(gateway: Gateway) -> None: + day: Final = _unique_day() + _seed( + day, + [ + (_row_id(), "user-a", day, "key-a", "gpt-5", "", "bedrock", None, "/v1/chat/completions", 10, 1.0, 1, 1), + (_row_id(), "user-b", day, "key-b", "gpt-5", "", "bedrock", None, "/v1/chat/completions", 10, 2.0, 1, 1), + ], + ) + _seed_guardrail_units( + day, + [ + ("pii-filter", "key-a", "contentPolicyUnits", 0.10, 5, 2), + ("pii-filter", "key-a", "sensitiveInformationPolicyUnits", 0.05, 1, 0), + ("pii-filter", "key-b", "contentPolicyUnits", 0.20, 1, 0), + ("pii-filter", "key-b", "legacyUnits", None, 7, 0), + ("topic-filter", "key-a", "topicPolicyUnits", 0.01, 1, 0), + ("topic-filter", "key-b", "topicPolicyUnits", None, 1, 0), + ("pii-filter", "orphan-key", "contentPolicyUnits", 9.0, 1, 0), + ], + ) + try: + everyone: Final = _day_breakdown(gateway, day) + assert _guardrail_spend(everyone) == { + "pii-filter": (pytest.approx(0.35), {"key-a": pytest.approx(0.15), "key-b": pytest.approx(0.20)}), + "topic-filter": (pytest.approx(0.01), {"key-a": pytest.approx(0.01), "key-b": 0.0}), + } + api_keys: Final = object_value(everyone["api_keys"]) + pii_filter: Final = object_value(object_value(everyone["guardrails"])["pii-filter"]) + assert object_value(pii_filter["metrics"]) == {"spend": pytest.approx(0.35), "untracked_units": 9} + pii_api_keys: Final = object_value(pii_filter["api_key_breakdown"]) + assert object_value(object_value(pii_api_keys["key-a"])["metrics"]) == { + "spend": pytest.approx(0.15), + "untracked_units": 2, + } + assert object_value(object_value(pii_api_keys["key-b"])["metrics"]) == { + "spend": pytest.approx(0.20), + "untracked_units": 7, + } + + assert object_value(everyone["guardrails"])["topic-filter"] == { + "metrics": {"spend": pytest.approx(0.01), "untracked_units": 1}, + "metadata": {"guardrail_name": "topic-filter"}, + "api_key_breakdown": { + "key-a": { + "metrics": {"spend": pytest.approx(0.01), "untracked_units": 0}, + "metadata": object_value(api_keys["key-a"])["metadata"], + }, + "key-b": { + "metrics": {"spend": 0.0, "untracked_units": 1}, + "metadata": object_value(api_keys["key-b"])["metadata"], + }, + }, + } + assert object_value(object_value(object_value(everyone["models"])["gpt-5"])["metrics"])["spend"] == 3.0 + + assert _guardrail_spend(_day_breakdown(gateway, day, user_id="user-a")) == { + "pii-filter": (pytest.approx(0.15), {"key-a": pytest.approx(0.15)}), + "topic-filter": (pytest.approx(0.01), {"key-a": pytest.approx(0.01)}), + } + assert _guardrail_spend(_day_breakdown(gateway, day, api_key_limit="1")) == { + "pii-filter": (pytest.approx(0.35), {"key-b": pytest.approx(0.20)}), + "topic-filter": (pytest.approx(0.01), {"key-b": 0.0}), + } + limited_pii_filter: Final = object_value( + object_value(_day_breakdown(gateway, day, api_key_limit="1")["guardrails"])["pii-filter"] + ) + assert object_value(limited_pii_filter["metrics"])["untracked_units"] == 9 + assert set(object_value(limited_pii_filter["api_key_breakdown"])) == {"key-b"} + assert _guardrail_spend(_day_breakdown(gateway, day, model="gpt-5")) == {} + assert _guardrail_spend(_day_breakdown(gateway, day, model="")) == _guardrail_spend(everyone) + finally: + _clean(day) + write_rows('DELETE FROM "LiteLLM_DailyGuardrailUsageUnits" WHERE date = %s', (day,)) + + +def _entity_breakdown(gateway: Gateway, url: str, day: str, **params: str) -> dict[str, JsonValue]: + response: Final = gateway.request("GET", url, params={"start_date": day, "end_date": day, **params}) + assert response.status_code == 200, response.text + results: Final = _RESULTS.validate_python(object_value(response.json())["results"]) + assert len(results) == 1 + return object_value(object_value(results[0])["breakdown"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("url", "table", "entity_column", "filter_param"), + [ + ("/team/daily/activity/aggregated", "LiteLLM_DailyTeamSpend", "team_id", "team_ids"), + ( + "/organization/daily/activity/aggregated", + "LiteLLM_DailyOrganizationSpend", + "organization_id", + "organization_ids", + ), + ], +) +async def test_team_and_organization_activity_break_down_guardrail_spend_by_name( + gateway: Gateway, url: str, table: str, entity_column: str, filter_param: str +) -> None: + day: Final = _unique_day() + db_guardrail_id: Final = str(uuid.uuid4()) + db_guardrail_name: Final = f"pii-filter-{db_guardrail_id[:8]}" + for entity_id, api_key, spend in (("entity-a", "key-a", 1.0), ("entity-b", "key-b", 2.0)): + write_rows( + f'INSERT INTO "{table}" (id, {entity_column}, date, api_key, model, spend, api_requests,' + " successful_requests, updated_at) VALUES (%s, %s, %s, %s, 'gpt-5', %s, 1, 1, now())", + (_row_id(), entity_id, day, api_key, str(spend)), + ) + write_rows( + 'INSERT INTO "LiteLLM_GuardrailsTable" (guardrail_id, guardrail_name, litellm_params, updated_at)' + " VALUES (%s, %s, '{}', now())", + (db_guardrail_id, db_guardrail_name), + ) + _seed_guardrail_units( + day, + [ + (db_guardrail_id, "key-a", "contentPolicyUnits", 0.10, 1, 0), + (db_guardrail_id, "key-b", "contentPolicyUnits", 0.20, 1, 0), + ("config-topic-filter", "key-a", "topicPolicyUnits", 0.05, 1, 0), + (db_guardrail_id, "orphan-key", "contentPolicyUnits", 9.0, 1, 0), + ], + ) + try: + everyone: Final = _entity_breakdown(gateway, url, day) + assert _guardrail_spend(everyone) == { + db_guardrail_id: (pytest.approx(0.30), {"key-a": pytest.approx(0.10), "key-b": pytest.approx(0.20)}), + "config-topic-filter": (pytest.approx(0.05), {"key-a": pytest.approx(0.05)}), + } + guardrails: Final = object_value(everyone["guardrails"]) + assert object_value(guardrails[db_guardrail_id])["metadata"] == {"guardrail_name": db_guardrail_name} + assert object_value(guardrails["config-topic-filter"])["metadata"] == {"guardrail_name": "config-topic-filter"} + + assert _guardrail_spend(_entity_breakdown(gateway, url, day, **{filter_param: "entity-a"})) == { + db_guardrail_id: (pytest.approx(0.10), {"key-a": pytest.approx(0.10)}), + "config-topic-filter": (pytest.approx(0.05), {"key-a": pytest.approx(0.05)}), + } + finally: + write_rows(f'DELETE FROM "{table}" WHERE date = %s', (day,)) + write_rows('DELETE FROM "LiteLLM_DailyGuardrailUsageUnits" WHERE date = %s', (day,)) + write_rows('DELETE FROM "LiteLLM_GuardrailsTable" WHERE guardrail_id = %s', (db_guardrail_id,)) diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 06fbcf58591..8bcb48a09b5 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,6 +1,4 @@ from collections.abc import AsyncIterator, Generator -from importlib import import_module -from typing import Final import pytest import pytest_asyncio @@ -23,16 +21,6 @@ def recording_server() -> Generator[RecordingServer]: yield server -def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: - try: - import_module("litellm.rust_bridge._native") - except ImportError: - skip: Final = pytest.mark.skip(reason="requires a compiled litellm.rust_bridge._native extension") - for item in items: - if "test_litellm_rust" in item.path.parts: - item.add_marker(skip) - - @pytest.fixture def isolated_azure_auth(monkeypatch: pytest.MonkeyPatch) -> None: for name in ( diff --git a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py index 09f09170241..cd3f7622fa2 100644 --- a/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py +++ b/tests/unit/a2a_protocol/providers/bedrock_agentcore/test_bedrock_agentcore_a2a.py @@ -10,6 +10,7 @@ Verifies that: """ import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -560,6 +561,37 @@ class TestStreaming: assert sent_body["params"]["message"]["messageId"] == "msg-001" assert [event["result"]["kind"] for event in events] == ["task", "status-update"] + @pytest.mark.asyncio + @pytest.mark.usefixtures("httpx_transport") + async def test_streaming_keeps_unicode_line_separators_inside_an_sse_event(self) -> None: + from litellm.a2a_protocol.providers.bedrock_agentcore.config import ( + BedrockAgentCoreA2AConfig, + ) + + event: Final = { + "jsonrpc": "2.0", + "id": "req-001", + "result": {"kind": "artifact-update", "artifact": {"parts": [{"kind": "text", "text": "a\u2028b\u2029c"}]}}, + } + with respx.mock(assert_all_called=True) as router: + router.post(url__regex=r".*/invocations.*").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(event, ensure_ascii=False)}\n\n", + ) + ) + events: Final = [ + streamed + async for streamed in BedrockAgentCoreA2AConfig().handle_streaming( + request_id="req-001", + params=SAMPLE_PARAMS, + litellm_params=SAMPLE_LITELLM_PARAMS, + ) + ] + + assert events == [event] + class TestConfigManager: """Test that config manager routes 'bedrock' correctly.""" diff --git a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py index a300560ae9d..ac3f0390f84 100644 --- a/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py +++ b/tests/unit/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -1,6 +1,7 @@ import asyncio import json import time +from typing import Final import httpx import pytest @@ -36,13 +37,8 @@ class _ShortTtlTokenClient: return _JsonResponse({"access_token": f"token-{self.calls}", "expires_in": 30}) -class _SSELines: - def __init__(self, lines): - self.lines = lines - - async def aiter_lines(self): - for line in self.lines: - yield line +def _sse_response(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) class _InvalidJsonStreamResponse: @@ -218,7 +214,7 @@ def test_cp4d_token_ttl_from_absolute_expiration(): @pytest.mark.asyncio async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): - response = _SSELines( + response = _sse_response( [ "data: null", "data: true", @@ -230,6 +226,14 @@ async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): ) +@pytest.mark.asyncio +async def test_accumulate_wxo_sse_text_keeps_unicode_line_separators_inside_an_event() -> None: + separated_text: Final = "before\u2028between\u2029after" + response: Final = _sse_response([f"data: {json.dumps({'results': separated_text}, ensure_ascii=False)}", ""]) + + assert await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response) == separated_text + + @pytest.mark.asyncio async def test_short_lived_tokens_are_not_served_from_cache(): client = _ShortTtlTokenClient() diff --git a/tests/unit/accounting/__init__.py b/tests/unit/accounting/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/accounting/test_measurements.py b/tests/unit/accounting/test_measurements.py new file mode 100644 index 00000000000..b65c3076c5d --- /dev/null +++ b/tests/unit/accounting/test_measurements.py @@ -0,0 +1,235 @@ +from datetime import datetime, timezone +from decimal import Decimal +from typing import Final + +import pytest + +from litellm.accounting import ( + CostEvent, + Dimension, + DimensionName, + InboundUsageEvent, + Measurement, + Metric, + MetricName, + Outcome, + OutboundUsageEvent, + PricingSnapshot, + measurements_for, +) +from litellm.types.utils import CallTypes + +_START: Final = datetime(2026, 1, 1, 23, 59, 59, tzinfo=timezone.utc) +_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc) +_DIMENSIONS: Final = ( + Dimension(name=DimensionName.TEAM, value="team"), + Dimension(name=DimensionName.TAG, value="alpha"), + Dimension(name=DimensionName.TAG, value="beta"), +) + + +def _snapshot(amount: Decimal) -> PricingSnapshot: + return PricingSnapshot( + calculator_version="test", + model="test-model", + provider=None, + pricing_source="upstream_header", + priced_model_key="test-model", + upstream_reported_cost=amount, + total_cost=amount, + ) + + +@pytest.mark.parametrize( + ("outcome", "outcome_metric"), + [ + (Outcome.SUCCEEDED, MetricName.SUCCESSFUL_REQUESTS), + (Outcome.FAILED, MetricName.FAILED_REQUESTS), + (Outcome.CANCELLED, MetricName.FAILED_REQUESTS), + ], +) +def test_caller_measurements_keep_both_request_and_outcome_counts(outcome, outcome_metric): + event = InboundUsageEvent( + event_id="caller-event", + request_id="caller", + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=outcome, + dimensions=_DIMENSIONS, + metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)), Metric(name=outcome_metric, value=Decimal(1))), + ) + assert measurements_for(event) == ( + Measurement( + source_event_id="caller-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.REQUESTS, value=Decimal(1)), + ), + Measurement( + source_event_id="caller-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=outcome_metric, value=Decimal(1)), + ), + ) + + +def test_usage_measurements_keep_every_metric_in_the_start_time_bucket(): + event = OutboundUsageEvent( + event_id="usage-event", + attempt_id="attempt", + request_id="caller", + retry_of_attempt_id=None, + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.SUCCEEDED, + dimensions=_DIMENSIONS, + metrics=( + Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)), + Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), + Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)), + Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)), + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)), + ), + ) + assert measurements_for(event) == ( + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)), + ), + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(2)), + ), + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)), + ), + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.CACHE_READ_INPUT_TOKENS, value=Decimal(3)), + ), + Measurement( + source_event_id="usage-event", + occurred_at=_START, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(1)), + ), + ) + + +@pytest.mark.parametrize("metrics", [None, ()]) +def test_unknown_or_empty_usage_does_not_fabricate_zero_measurements(metrics): + event = OutboundUsageEvent( + event_id="usage-event", + attempt_id="attempt", + request_id=None, + retry_of_attempt_id=None, + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.SUCCEEDED, + dimensions=(), + metrics=metrics, + ) + assert measurements_for(event) == () + + +@pytest.mark.parametrize("amount", [Decimal(0), Decimal("0.1234567890123456789")]) +def test_cost_measurements_preserve_authoritative_amount_identity_and_occurrence(amount): + event = CostEvent( + event_id="cost-event", + usage_event_id="usage-event", + occurred_at=_END, + dimensions=_DIMENSIONS, + charges=(Metric(name=MetricName.SPEND, value=amount, unit="USD"),), + snapshot=_snapshot(amount), + ) + assert measurements_for(event) == ( + Measurement( + source_event_id="cost-event", + occurred_at=_END, + dimensions=_DIMENSIONS, + metric=Metric(name=MetricName.SPEND, value=amount, unit="USD"), + ), + ) + + +def test_billed_retry_and_success_contribute_one_caller_request_and_both_charges(): + caller = InboundUsageEvent( + event_id="caller-event", + request_id="caller", + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.SUCCEEDED, + dimensions=_DIMENSIONS, + metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), + ) + retry = OutboundUsageEvent( + event_id="retry-usage", + attempt_id="retry", + request_id="caller", + retry_of_attempt_id=None, + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.FAILED, + dimensions=_DIMENSIONS, + metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(4)),), + ) + success = OutboundUsageEvent( + event_id="success-usage", + attempt_id="success", + request_id="caller", + retry_of_attempt_id="retry", + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_END, + outcome=Outcome.SUCCEEDED, + dimensions=_DIMENSIONS, + metrics=(Metric(name=MetricName.TOTAL_TOKENS, value=Decimal(12)),), + ) + retry_cost = CostEvent( + event_id="retry-cost", + usage_event_id="retry-usage", + occurred_at=_START, + dimensions=_DIMENSIONS, + charges=(Metric(name=MetricName.SPEND, value=Decimal("0.025"), unit="USD"),), + snapshot=_snapshot(Decimal("0.025")), + ) + success_cost = CostEvent( + event_id="success-cost", + usage_event_id="success-usage", + occurred_at=_START, + dimensions=_DIMENSIONS, + charges=(Metric(name=MetricName.SPEND, value=Decimal("0.125"), unit="USD"),), + snapshot=_snapshot(Decimal("0.125")), + ) + measurements = ( + *measurements_for(caller), + *measurements_for(retry), + *measurements_for(success), + *measurements_for(retry_cost), + *measurements_for(success_cost), + ) + assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.REQUESTS) == Decimal(1) + assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.TOTAL_TOKENS) == Decimal(16) + assert sum(m.metric.value for m in measurements if m.metric.name is MetricName.SPEND) == Decimal("0.150") + assert {m.source_event_id for m in measurements} == { + "caller-event", + "retry-usage", + "success-usage", + "retry-cost", + "success-cost", + } diff --git a/tests/unit/accounting/test_types.py b/tests/unit/accounting/test_types.py new file mode 100644 index 00000000000..f645ee6c2cc --- /dev/null +++ b/tests/unit/accounting/test_types.py @@ -0,0 +1,484 @@ +import re +from collections.abc import Mapping +from datetime import datetime, timedelta, timezone +from decimal import Decimal +from enum import Enum +from pathlib import Path +from typing import Final + +import pytest +from pydantic import ValidationError + +from litellm.accounting import ( + AdmissionDecision, + AggregateQuery, + AggregateRow, + CostEvent, + DeliveryState, + DeliveryTarget, + Dimension, + DimensionFilter, + DimensionName, + EventKind, + Granularity, + InboundUsageEvent, + Metric, + MetricName, + Multiplier, + Outcome, + OutboundUsageEvent, + PricedQuantity, + PricingSnapshot, + Rate, + validate_aggregate_rows, +) +from litellm.accounting.types import PricingUnit +from litellm.types.utils import CallTypes + +_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +_END: Final = datetime(2026, 1, 2, tzinfo=timezone.utc) +_REQUEST: Final = Metric(name=MetricName.REQUESTS, value=Decimal(1)) +_TOKEN: Final = Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(10)) +_SPEND: Final = Metric(name=MetricName.SPEND, value=Decimal("0.1234567890123456789"), unit="USD") +_TEAM: Final = Dimension(name=DimensionName.TEAM, value="team") + + +def _inbound(metrics=(_REQUEST,), outcome=Outcome.SUCCEEDED): + return InboundUsageEvent( + event_id="inbound:request", + request_id="request", + call_type=CallTypes.acompletion, + started_at=_START, + ended_at=_START, + outcome=outcome, + dimensions=(_TEAM,), + metrics=metrics, + ) + + +def _outbound(metrics=(_TOKEN,), **changes): + return OutboundUsageEvent.model_validate( + { + "event_id": "outbound:attempt", + "attempt_id": "attempt", + "request_id": "request", + "retry_of_attempt_id": None, + "call_type": CallTypes.acompletion, + "started_at": _START, + "ended_at": _END, + "outcome": Outcome.SUCCEEDED, + "dimensions": (_TEAM,), + "metrics": metrics, + **changes, + } + ) + + +def _query(**changes): + return AggregateQuery.model_validate( + { + "metric": MetricName.SPEND, + "filters": (DimensionFilter(name=DimensionName.TEAM, values=("one", "two")),), + "group_by": (DimensionName.MODEL,), + "start": _START, + "end": _END, + "granularity": Granularity.DAY, + **changes, + } + ) + + +def test_wire_values_preserve_decimal_precision_and_normalize_aware_times(): + offset = timezone(timedelta(hours=-8)) + event = CostEvent( + event_id="cost:attempt", + usage_event_id="outbound:attempt", + occurred_at=datetime(2025, 12, 31, 16, tzinfo=offset), + dimensions=(_TEAM,), + charges=(_SPEND,), + snapshot=_snapshot(), + ) + wire = event.model_dump(mode="json") + assert wire["occurred_at"] == "2026-01-01T00:00:00Z" + assert wire["charges"] == [{"name": "spend", "value": "0.1234567890123456789", "unit": "USD"}] + assert event.occurred_at == _START + assert CostEvent.model_validate_json(event.model_dump_json()).charges[0].value == _SPEND.value + + +@pytest.mark.parametrize( + ("name", "value", "unit", "message"), + [ + (MetricName.SPEND, Decimal(1), "count", "SPEND requires USD"), + (MetricName.REQUESTS, Decimal(1), "USD", "SPEND requires USD"), + (MetricName.PROMPT_TOKENS, Decimal("1.5"), "count", "integral"), + (MetricName.SPEND, Decimal("NaN"), "USD", "finite"), + (MetricName.SPEND, Decimal("Infinity"), "USD", "finite"), + (MetricName.SPEND, Decimal(-1), "USD", "greater than or equal"), + (MetricName.SPEND, 0.1, "USD", "instance of Decimal"), + ], +) +def test_metrics_reject_wrong_units_nonfinite_values_and_binary_float_inputs(name, value, unit, message): + with pytest.raises(ValidationError, match=message): + Metric(name=name, value=value, unit=unit) + + +@pytest.mark.parametrize("metrics", [(_TOKEN,), (_SPEND,), ()]) +def test_inbound_activity_cannot_own_usage_or_cost_or_omit_request_count(metrics): + with pytest.raises(ValidationError, match="Inbound events"): + _inbound(metrics) + + +@pytest.mark.parametrize( + ("outcome", "metric"), + [ + (Outcome.FAILED, Metric(name=MetricName.SUCCESSFUL_REQUESTS, value=Decimal(1))), + (Outcome.SUCCEEDED, Metric(name=MetricName.FAILED_REQUESTS, value=Decimal(1))), + ], +) +def test_inbound_outcome_counts_cannot_contradict_outcome(outcome, metric): + with pytest.raises(ValidationError, match="count must match outcome"): + _inbound((_REQUEST, metric), outcome) + + +@pytest.mark.parametrize("metrics", [(_REQUEST,), (_SPEND,)]) +def test_outbound_attempts_cannot_own_caller_counts_or_charges(metrics): + with pytest.raises(ValidationError, match="Outbound events own only"): + _outbound(metrics) + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"started_at": datetime(2026, 1, 1)}, "timezone"), + ({"ended_at": _START - timedelta(seconds=1)}, "End must follow start"), + ({"retry_of_attempt_id": "attempt"}, "cannot retry itself"), + ({"dimensions": (_TEAM, _TEAM)}, "Duplicate dimension values"), + ({"dimensions": (_TEAM, Dimension(name=DimensionName.TEAM, value="other"))}, "Only tag and model access group"), + ( + { + "metrics": ( + Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(5)), + Metric(name=MetricName.REASONING_TOKENS, value=Decimal(6)), + ) + }, + "reasoning_tokens cannot exceed completion_tokens", + ), + ( + { + "metrics": ( + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(5)), + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(6)), + ) + }, + "cache_creation_input_tokens_above_1hr cannot exceed cache_creation_input_tokens", + ), + ({"event_id": " "}, "pattern"), + ({"metrics": (_TOKEN, _TOKEN)}, "Duplicate metric names"), + ], +) +def test_outbound_events_reject_invalid_identity_time_and_duplicate_attribution(changes, message): + with pytest.raises(ValidationError, match=message): + _outbound(**changes) + + +def test_unknown_and_known_zero_usage_are_distinct_and_values_are_immutable(): + unknown = _outbound(None) + zero = _outbound((Metric(name=MetricName.PROMPT_TOKENS, value=Decimal(0)),)) + assert unknown.metrics is None + assert zero.metrics is not None and zero.metrics[0].value == 0 + assert _inbound().started_at == _inbound().ended_at + with pytest.raises(ValidationError, match="frozen"): + zero.attempt_id = "different" + with pytest.raises(ValidationError, match="frozen"): + zero.metrics[0].value = Decimal(10) + + +@pytest.mark.parametrize("charges", [(), (_TOKEN,), (_SPEND, _SPEND)]) +def test_cost_events_require_one_authoritative_spend(charges): + with pytest.raises(ValidationError, match=r"authoritative USD spend|Duplicate metric"): + CostEvent( + event_id="cost", + usage_event_id="usage", + occurred_at=_START, + dimensions=(), + charges=charges, + snapshot=_snapshot(), + ) + + +def test_cost_event_cannot_reuse_its_usage_identity(): + with pytest.raises(ValidationError, match="distinct identities"): + CostEvent( + event_id="usage", + usage_event_id="usage", + occurred_at=_START, + dimensions=(), + charges=(_SPEND,), + snapshot=_snapshot(), + ) + + +@pytest.mark.parametrize("granularity", [Granularity.DAY, None]) +def test_daily_and_whole_interval_queries_preserve_filter_and_group_semantics(granularity): + query = _query(granularity=granularity) + assert query.granularity is granularity + assert query.filters[0].values == ("one", "two") + assert query.group_by == (DimensionName.MODEL,) + assert query.start == _START and query.end == _END + + +@pytest.mark.parametrize( + ("changes", "message"), + [ + ({"end": _START}, "End must follow start"), + ({"start": _START + timedelta(hours=1)}, "UTC day boundaries"), + ({"end": _END + timedelta(microseconds=1)}, "UTC day boundaries"), + ({"granularity": "hour"}, "day"), + ({"group_by": (DimensionName.TEAM, DimensionName.TEAM)}, "Duplicate group-by"), + ({"filters": (DimensionFilter(name=DimensionName.TEAM, values=("one",)),) * 2}, "one IN filter"), + ({"filters": ({"name": "team", "values": ()},)}, "at least 1"), + ], +) +def test_queries_reject_unsupported_or_ambiguous_ranges_and_filters(changes, message): + with pytest.raises(ValidationError, match=message): + _query(**changes) + + +def test_rows_allow_adjacent_buckets_and_overlapping_distinct_groups(): + first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM,), value=Decimal(1)) + second = AggregateRow(start=_END, end=_END + timedelta(days=1), dimensions=(_TEAM,), value=Decimal(2)) + other = AggregateRow(start=_START, end=_END, dimensions=(), value=Decimal(3)) + rows = (second, other, first) + assert validate_aggregate_rows(rows) is rows + assert validate_aggregate_rows(()) == () + + +def test_rows_reject_overlap_within_a_group_regardless_of_dimension_order(): + model = Dimension(name=DimensionName.MODEL, value="model") + first = AggregateRow(start=_START, end=_END, dimensions=(_TEAM, model), value=Decimal(1)) + overlapping = AggregateRow(start=_START, end=_END, dimensions=(model, _TEAM), value=Decimal(2)) + with pytest.raises(ValueError, match="overlap within a dimension group"): + validate_aggregate_rows((first, overlapping)) + + +def test_rows_reject_empty_intervals(): + with pytest.raises(ValidationError, match="End must follow start"): + AggregateRow(start=_START, end=_START, dimensions=(), value=Decimal(0)) + + +def test_admission_decision_preserves_a_denial_reason(): + decision = AdmissionDecision(allowed=False, reason="Team budget exhausted") + assert not decision.allowed + assert decision.reason == "Team budget exhausted" + + +def _snapshot(**changes: object) -> PricingSnapshot: + return PricingSnapshot.model_validate( + { + "calculator_version": "1.80.0", + "model": "gpt-4o", + "provider": "openai", + "pricing_source": "cost_map", + "priced_model_key": "openai/gpt-4o", + "multipliers": (Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")),), + "rates": (Rate(measure="input", unit="token", usd_per_unit=Decimal("0.0000025")),), + "usage": (PricedQuantity(measure="input", unit="token", value=_TOKEN.value),), + "input_cost": Decimal("0.000025"), + "total_cost": _SPEND.value, + **changes, + } + ) + + +def test_cost_event_carries_a_snapshot_whose_total_matches_the_spend(): + event = CostEvent( + event_id="cost:attempt", + usage_event_id="outbound:attempt", + occurred_at=_START, + dimensions=(_TEAM,), + charges=(_SPEND,), + snapshot=_snapshot(), + ) + assert event.snapshot is not None + assert event.snapshot.total_cost == _SPEND.value + assert event.snapshot.multipliers[0].factor == Decimal("1.1") + with pytest.raises(ValidationError, match="Snapshot total must equal the authoritative spend"): + CostEvent( + event_id="cost:attempt", + usage_event_id="outbound:attempt", + occurred_at=_START, + dimensions=(_TEAM,), + charges=(_SPEND,), + snapshot=_snapshot(total_cost=Decimal("0.5")), + ) + + +@pytest.mark.parametrize( + ("changes", "message"), + ( + ( + { + "rates": ( + Rate(measure="input", unit="token", usd_per_unit=Decimal(1)), + Rate(measure="input", unit="token", usd_per_unit=Decimal(2)), + ) + }, + "Duplicate rate measures and units", + ), + ({"usage": (_SPEND,)}, "PricedQuantity"), + ({"usage": (_REQUEST,)}, "PricedQuantity"), + ({"pricing_source": "guess"}, "pricing_source"), + ({"multipliers": ({"name": "x", "value": "y", "factor": Decimal(0)},)}, "factor"), + ({"schema_version": 2}, "schema_version"), + ( + {"additional_costs": ({"name": "fee", "amount": Decimal(1)}, {"name": "fee", "amount": Decimal(2)})}, + "Duplicate additional cost names", + ), + ), +) +def test_snapshots_reject_ambiguous_rates_wrong_usage_and_unknown_sources(changes: Mapping[str, object], message: str): + with pytest.raises(ValidationError, match=message): + _snapshot(**changes) + + +@pytest.mark.parametrize( + ("unit", "value"), + [("token", "10"), ("second", "1.25"), ("character", "25"), ("image", "2"), ("request", "1")], +) +def test_snapshot_preserves_matching_rates_and_quantities_for_each_billing_unit(unit: PricingUnit, value: str) -> None: + quantity: Final = PricedQuantity(measure="input", unit=unit, value=Decimal(value)) + rate: Final = Rate(measure="input", unit=unit, usd_per_unit=Decimal("0.003")) + amount: Final = quantity.value * rate.usd_per_unit + snapshot: Final = _snapshot(rates=(rate,), usage=(quantity,), total_cost=amount) + restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json()) + assert restored == snapshot + assert restored.usage[0].value * restored.rates[0].usd_per_unit == amount + + +@pytest.mark.parametrize( + ("usage", "message"), + [ + ((), "Every rate requires"), + (({"measure": "output", "unit": "token", "value": Decimal(10)},), "Every rate requires"), + (({"measure": "input", "unit": "second", "value": Decimal(10)},), "Every rate requires"), + ( + ({"measure": "input", "unit": "token", "value": Decimal(10)},) * 2, + "Duplicate priced quantity", + ), + ], +) +def test_snapshot_rejects_missing_mismatched_or_duplicate_priced_quantities( + usage: tuple[Mapping[str, object], ...], message: str +) -> None: + with pytest.raises(ValidationError, match=message): + _snapshot(usage=usage) + + +@pytest.mark.parametrize("unit", ["token", "character", "image", "request"]) +def test_priced_counts_remain_integral(unit: PricingUnit) -> None: + with pytest.raises(ValidationError, match="integral"): + PricedQuantity(measure="input", unit=unit, value=Decimal("1.25")) + + +@pytest.mark.parametrize( + ("value", "message"), + [ + (Decimal(-1), "greater than or equal"), + (Decimal("NaN"), "finite"), + (Decimal("Infinity"), "finite"), + (1.25, "instance of Decimal"), + ], +) +def test_pricing_values_reject_invalid_amounts(value: Decimal | float, message: str) -> None: + with pytest.raises(ValidationError, match=message): + PricedQuantity.model_validate({"measure": "input", "unit": "second", "value": value}) + with pytest.raises(ValidationError, match=message): + Rate.model_validate({"measure": "input", "unit": "second", "usd_per_unit": value}) + + +@pytest.mark.parametrize("snapshot_fields", [{}, {"snapshot": None}]) +def test_cost_event_cannot_omit_its_derivation(snapshot_fields: Mapping[str, object]) -> None: + with pytest.raises(ValidationError, match="snapshot"): + CostEvent.model_validate( + { + "event_id": "cost", + "usage_event_id": "usage", + "occurred_at": _START, + "dimensions": (), + "charges": (_SPEND,), + **snapshot_fields, + } + ) + + +def test_upstream_charge_snapshot_preserves_reported_cost_without_inventing_rates() -> None: + snapshot: Final = _snapshot( + pricing_source="upstream_header", rates=(), usage=(), upstream_reported_cost=_SPEND.value + ) + restored: Final = PricingSnapshot.model_validate_json(snapshot.model_dump_json()) + assert restored.upstream_reported_cost == restored.total_cost == _SPEND.value + assert restored.rates == restored.usage == () + + +def test_subset_token_counts_round_trip_without_changing_totals() -> None: + metrics: Final = ( + Metric(name=MetricName.COMPLETION_TOKENS, value=Decimal(10)), + Metric(name=MetricName.REASONING_TOKENS, value=Decimal(3)), + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS, value=Decimal(8)), + Metric(name=MetricName.CACHE_CREATION_INPUT_TOKENS_ABOVE_1HR, value=Decimal(2)), + ) + event: Final = _outbound(metrics=metrics) + assert OutboundUsageEvent.model_validate_json(event.model_dump_json()).metrics == metrics + + +@pytest.mark.parametrize( + "schema_path", + [ + "schema.prisma", + "litellm/proxy/schema.prisma", + "litellm-proxy-extras/litellm_proxy_extras/schema.prisma", + ], +) +@pytest.mark.parametrize( + ("enum_name", "python_enum"), + [ + ("AccountingEventKind", EventKind), + ("AccountingDeliveryState", DeliveryState), + ("AccountingDeliveryTarget", DeliveryTarget), + ], +) +def test_prisma_accepts_exactly_the_public_accounting_enum_values( + schema_path: str, enum_name: str, python_enum: type[Enum] +) -> None: + schema: Final = (Path(__file__).resolve().parents[3] / schema_path).read_text() + declaration: Final = re.search(r"enum " + enum_name + r"\s*\{([^}]+)\}", schema) + assert declaration is not None, f"Missing persisted enum {enum_name}" + assert frozenset(declaration[1].split()) == frozenset(member.value for member in python_enum) + + +def test_one_event_carries_every_matched_model_access_group(): + groups = ( + Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="gpt-group"), + Dimension(name=DimensionName.MODEL_ACCESS_GROUP, value="claude-group"), + ) + assert _outbound(dimensions=(_TEAM, *groups)).dimensions == (_TEAM, *groups) + with pytest.raises(ValidationError, match="Duplicate dimension values"): + _outbound(dimensions=(groups[0], groups[0])) + + +def test_snapshots_reject_duplicate_multipliers_and_contradictory_token_counts(): + multiplier = Multiplier(name="data_residency", value="eu", factor=Decimal("1.1")) + with pytest.raises(ValidationError, match="Duplicate multiplier names"): + _snapshot(multipliers=(multiplier, multiplier)) + with pytest.raises(ValidationError, match="reasoning_tokens cannot exceed completion_tokens"): + _snapshot( + usage=( + PricedQuantity(measure="input", unit="token", value=_TOKEN.value), + PricedQuantity(measure="completion_tokens", unit="token", value=Decimal(5)), + PricedQuantity(measure="reasoning_tokens", unit="token", value=Decimal(6)), + ) + ) + + +def test_agent_live_spend_is_a_delivery_target(): + assert DeliveryTarget.AGENT_SPEND.value == "LiteLLM_AgentsTable" diff --git a/tests/unit/caching/test_in_memory_cache.py b/tests/unit/caching/test_in_memory_cache.py index 75232476ecd..901832937b5 100644 --- a/tests/unit/caching/test_in_memory_cache.py +++ b/tests/unit/caching/test_in_memory_cache.py @@ -235,13 +235,9 @@ def test_in_memory_cache_eviction_order(): # Add items with different TTLs now = time.time() - in_memory_cache.set_cache( - key="early_expire", value="value_1", ttl=100 - ) # expires in 100 seconds + in_memory_cache.set_cache(key="early_expire", value="value_1", ttl=100) # expires in 100 seconds time.sleep(0.01) - in_memory_cache.set_cache( - key="late_expire", value="value_2", ttl=200 - ) # expires in 200 seconds + in_memory_cache.set_cache(key="late_expire", value="value_2", ttl=200) # expires in 200 seconds # Verify TTL order early_ttl = in_memory_cache.ttl_dict["early_expire"] @@ -379,3 +375,20 @@ def test_in_memory_cache_injected_clock_controls_expiry_and_eviction() -> None: assert cache.get_cache("second") is None assert cache.get_cache("third") == "replacement" assert cache.get_cache("fourth") == "new" + + +@pytest.mark.asyncio +async def test_async_increment_pipeline_keeps_each_key_for_its_own_ttl() -> None: + now: Final = {"t": 1_000.0} + cache: Final = InMemoryCache(default_ttl=600, clock=lambda: now["t"]) + + await cache.async_increment_pipeline( + increment_list=[ + {"key": "minute", "increment_value": 1, "ttl": 60}, + {"key": "day", "increment_value": 1, "ttl": 86400}, + ] + ) + now["t"] += 601 + + assert cache.get_cache("minute") is None + assert cache.get_cache("day") == 1 diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 9aa494c1aab..46f1dbf2e7a 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1102,6 +1102,24 @@ def test_recover_output_items_merges_text_only_items_at_distinct_indices(): assert recovered[1]["content"][0]["text"] == "From OUTPUT_TEXT_DONE only" +def test_recover_output_items_keeps_an_item_whose_text_holds_unicode_line_separators() -> None: + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + item: Final = { + "type": "message", + "id": "msg_item_0", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "From\u2028OUTPUT_ITEM\u2029DONE", "annotations": []}], + } + event: Final = {"type": "response.output_item.done", "output_index": 0, "item": item} + raw_sse: Final = f"data: {json.dumps(event, ensure_ascii=False)}\ndata: [DONE]\n" + + assert LiteLLMResponsesTransformationHandler._recover_output_items_from_raw_sse(raw_sse) == [item] + + def test_transform_response_prefers_completed_output_from_raw_sse(): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index b9813289749..e5501310556 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -578,8 +578,14 @@ async def test_unknown_provider_is_rejected_before_http(respx_mock: respx.MockRo @pytest.mark.asyncio -async def test_provider_without_decisions_support_is_rejected_before_http(respx_mock: respx.MockRouter) -> None: - with pytest.raises(litellm.BadRequestError, match=r"Unknown Decisions provider 'anthropic'\. Supported providers"): +async def test_provider_without_decisions_support_is_rejected_before_http( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + with pytest.raises( + litellm.BadRequestError, match=r"Unknown Decisions provider 'anthropic'\. Supported providers" + ) as error: await litellm.asystemone( model="anthropic/claude-sonnet-4-5", state="review", @@ -587,6 +593,8 @@ async def test_provider_without_decisions_support_is_rejected_before_http(respx_ api_key="caller-key", ) + supported: Final = str(error.value).split("Supported providers: ", 1)[1].split(", ") + assert {"ollama", "ollama_chat"} <= set(supported) assert len(respx_mock.calls) == 0 @@ -1536,6 +1544,8 @@ _PROVIDERS_WHOSE_CHAT_MODELS_ALSO_DECIDE: Final = frozenset( LlmProviders.OPENROUTER.value, LlmProviders.HOSTED_VLLM.value, LlmProviders.LM_STUDIO.value, + LlmProviders.OLLAMA.value, + LlmProviders.OLLAMA_CHAT.value, } ) _DEDICATED_DECIDER_PROVIDERS: Final[tuple[str, ...]] = tuple( @@ -3911,6 +3921,185 @@ async def test_native_decisions_distinguishes_missing_usage_from_reported_zero( assert (expected == 0) == (usage["input_tokens"] == usage["output_tokens"] == 0) +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("use_async", (True, False), ids=("asystemone", "systemone")) +@pytest.mark.parametrize("api_base", ("http://ollama.local:11434", "http://ollama.local:11434/v1")) +@pytest.mark.parametrize("api_key", (None, "fake-api-key")) +async def test_ollama_systemone_preserves_questions_without_unrelated_or_placeholder_auth( + provider: str, + use_async: bool, + api_base: str, + api_key: str | None, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-openai-key") + monkeypatch.setattr(litellm, "api_key", "unrelated-global-key") + monkeypatch.setattr(litellm, "openai_key", "unrelated-global-openai-key") + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + kwargs: Final = { + "model": f"{provider}/d3:4b", + "state": "review", + "questions": _QUESTIONS, + "api_base": api_base, + "api_key": api_key, + } + + response: Final = await litellm.asystemone(**kwargs) if use_async else litellm.systemone(**kwargs) + + assert json.loads(route.calls.last.request.content) == { + "model": "d3:4b", + "state": "review", + "questions": dict(_QUESTIONS), + } + assert "authorization" not in route.calls.last.request.headers + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + assert response.answers["sentiment"].choice == "positive" + assert response.answers["severity"].score == 1 + assert response.usage == DecisionsUsage(input_tokens=_INPUT_TOKENS, output_tokens=_OUTPUT_TOKENS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("use_async", (True, False), ids=("adecisions", "decisions")) +async def test_ollama_decisions_translates_openai_questions_and_answers( + provider: str, use_async: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + kwargs: Final = { + "model": f"{provider}/d3:4b", + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + "api_base": "http://ollama.local:11434/v1", + } + + response: Final = await litellm.adecisions(**kwargs) if use_async else litellm.decisions(**kwargs) + + assert json.loads(route.calls.last.request.content) == { + "model": "d3:4b", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert "authorization" not in route.calls.last.request.headers + assert response.answers[0].type == "predicate" + assert response.answers[0].name == "is_defect" + assert response.answers[0].probability == 0.9 + assert response.usage is not None + assert response.usage.input_tokens == _INPUT_TOKENS + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("environment_base", (None, "http://ollama.local:11434/v1")) +async def test_ollama_decisions_resolves_environment_base_or_localhost_default( + provider: str, environment_base: str | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + if environment_base is not None: + monkeypatch.setenv("OLLAMA_API_BASE", environment_base) + url: Final = f"{environment_base or 'http://localhost:11434/v1'}/systemone" + route: Final = respx_mock.post(url).respond(json=_RESPONSE) + + response: Final = await litellm.asystemone(model=f"{provider}/d3:4b", state="review", questions=_QUESTIONS) + + assert route.call_count == 1 + assert "authorization" not in route.calls.last.request.headers + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("explicit_key", (None, "deployment-key")) +async def test_ollama_decisions_resolves_environment_key_and_explicit_connection_overrides( + provider: str, explicit_key: str | None, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + monkeypatch.setenv("OLLAMA_API_BASE", "http://ignored.local:11434") + monkeypatch.setenv("OLLAMA_API_KEY", "environment-key") + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + + response: Final = await litellm.asystemone( + model=f"{provider}/d3:4b", + state="review", + questions=_QUESTIONS, + api_base="http://ollama.local:11434", + api_key=explicit_key, + ) + + assert route.calls.last.request.headers["authorization"] == f"Bearer {explicit_key or 'environment-key'}" + assert response.answers["is_defect"] == NoulAnswer(type="noul", noul=0.9) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +async def test_ollama_decisions_maps_upstream_bad_requests( + provider: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond( + 400, json={"error": "decision rejected"} + ) + + with pytest.raises(litellm.BadRequestError, match="decision rejected") as error: + await litellm.asystemone( + model=f"{provider}/d3:4b", + state="review", + questions=_QUESTIONS, + api_base="http://ollama.local:11434", + num_retries=0, + ) + + assert error.value.status_code == 400 + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("drop_params", (False, True)) +@pytest.mark.parametrize("image", ("data:image/png;base64,AA==", "https://images.example/input.png")) +async def test_ollama_decisions_rejects_input_images_before_http_or_drops_them( + provider: str, drop_params: bool, image: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + monkeypatch.setattr(litellm, "drop_params", False) + kwargs: Final = { + "model": f"{provider}/d3:4b", + "input": [ + { + "role": "user", + "content": [{"type": "input_text", "text": "review"}, {"type": "input_image", "image_url": image}], + } + ], + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + "api_base": "http://ollama.local:11434", + "drop_params": drop_params, + } + if not drop_params: + with pytest.raises(litellm.UnsupportedParamsError, match=r"\['input_image'\].*drop_params"): + await litellm.adecisions(**kwargs) + assert len(respx_mock.calls) == 0 + return + route: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + + response: Final = await litellm.adecisions(**kwargs) + + assert json.loads(route.calls.last.request.content) == { + "model": "d3:4b", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert response.answers[0].probability == 0.9 + + @pytest.mark.asyncio @pytest.mark.parametrize("use_async", (True, False), ids=("asystemone", "systemone")) @pytest.mark.parametrize("api_base", ("http://lmstudio.local:1234", "http://lmstudio.local:1234/v1")) diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index 01f20689e89..f9f4e17c760 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -15,6 +15,7 @@ import pytest from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch +from litellm.llms.base_llm.files.transformation import FileListModelScope from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.openai_files_endpoints.common_utils import BATCH_CREATE_HIDDEN_PARAM from litellm.types.llms.openai import FileListPage, OpenAIFileObject @@ -649,6 +650,22 @@ async def test_afile_list_denies_a_caller_with_no_identity_at_all(): assert table.find_many_calls == [] +@pytest.mark.asyncio +async def test_afile_list_reads_no_rows_when_no_deployment_can_hold_a_file() -> None: + managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")]) + + response: Final = await managed_files.afile_list( + purpose=None, + litellm_parent_otel_span=None, + user_api_key_dict=_make_user_api_key_dict(), + model_scope=FileListModelScope(frozenset()), + ) + + assert response.data == [] + assert response.has_more is False + assert table.find_many_calls == [] + + @pytest.mark.asyncio async def test_afile_list_filters_by_purpose(): managed_files, _ = _make_managed_files_over_rows( diff --git a/tests/unit/google_genai/test_google_genai_streaming_iterator.py b/tests/unit/google_genai/test_google_genai_streaming_iterator.py index 120c99b6e42..cbb9638b3d9 100644 --- a/tests/unit/google_genai/test_google_genai_streaming_iterator.py +++ b/tests/unit/google_genai/test_google_genai_streaming_iterator.py @@ -62,7 +62,9 @@ def _logging_obj( ], ) def test_streaming_logging_targets_the_provider_that_served_the_request( - iterator_cls: type, + iterator_cls: ( + type[AsyncGoogleGenAIGenerateContentStreamingIterator] | type[GoogleGenAIGenerateContentStreamingIterator] + ), custom_llm_provider: str, expected_endpoint_type: EndpointType, ): @@ -99,18 +101,15 @@ def _large_inline_data_event() -> str: return f"data: {json.dumps(payload)}" +def _sse_response(body: str) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream(body.encode())) + + @pytest.mark.asyncio async def test_async_streaming_iterator_yields_complete_sse_events(): """Large inlineData must not be split across byte-chunk boundaries.""" - mock_response = MagicMock() - - async def _aiter_lines(): - yield _large_inline_data_event() - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(_large_inline_data_event() + "\n\n"), model="gemini-3.1-flash-image-preview", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -128,11 +127,8 @@ async def test_async_streaming_iterator_yields_complete_sse_events(): def test_sync_streaming_iterator_yields_complete_sse_events(): - mock_response = MagicMock() - mock_response.iter_lines.return_value = iter([_large_inline_data_event()]) - iterator = GoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(_large_inline_data_event() + "\n\n"), model="gemini-3.1-flash-image-preview", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -150,17 +146,8 @@ def test_sync_streaming_iterator_yields_complete_sse_events(): @pytest.mark.asyncio async def test_async_streaming_iterator_preserves_multi_field_sse_event(): - mock_response = MagicMock() - - async def _aiter_lines(): - yield "event: message" - yield 'data: {"text":"hi"}' - yield "" - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response('event: message\ndata: {"text":"hi"}\n\n'), model="gemini-test", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -174,16 +161,8 @@ async def test_async_streaming_iterator_preserves_multi_field_sse_event(): @pytest.mark.asyncio async def test_async_streaming_iterator_forwards_sse_comment_events(): - mock_response = MagicMock() - - async def _aiter_lines(): - yield ": keepalive" - yield "" - - mock_response.aiter_lines = _aiter_lines - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + response=_sse_response(": keepalive\n\n"), model="gemini-test", logging_obj=_logging_obj(), generate_content_provider_config=MagicMock(), @@ -195,6 +174,46 @@ async def test_async_streaming_iterator_forwards_sse_comment_events(): assert chunk == b": keepalive\n\n" +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _first_chunk( + iterator: GoogleGenAIGenerateContentStreamingIterator | AsyncGoogleGenAIGenerateContentStreamingIterator, +) -> bytes: + if isinstance(iterator, AsyncGoogleGenAIGenerateContentStreamingIterator): + return await iterator.__anext__() + return next(iterator) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "iterator_cls", + [AsyncGoogleGenAIGenerateContentStreamingIterator, GoogleGenAIGenerateContentStreamingIterator], +) +async def test_streaming_iterator_keeps_unicode_line_separators_inside_the_event( + iterator_cls: ( + type[AsyncGoogleGenAIGenerateContentStreamingIterator] | type[GoogleGenAIGenerateContentStreamingIterator] + ), +) -> None: + payload: Final = {"candidates": [{"content": {"parts": [{"text": _LINE_SEPARATED_TEXT}]}}]} + event: Final = f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" + iterator: Final = iterator_cls( + response=_sse_response(event), + model="gemini-3.8-flash", + logging_obj=_logging_obj(), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + chunk: Final = await _first_chunk(iterator) + + assert chunk == event.encode() + assert ( + json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0]["text"] == _LINE_SEPARATED_TEXT + ) + + def _gemini_event(text: str, finish_reason: str | None = None) -> str: candidate: Final = {"content": {"role": "model", "parts": [{"text": text}]}, "index": 0} payload: Final = { diff --git a/tests/unit/images/test_image_edit.py b/tests/unit/images/test_image_edit.py index fb7a2869b2a..4f858050e3e 100644 --- a/tests/unit/images/test_image_edit.py +++ b/tests/unit/images/test_image_edit.py @@ -105,6 +105,24 @@ def test_openai_image_edit_sync_accepts_bytesio_images(respx_mock: respx.MockRou assert _multipart_image_parts(route.calls[0].request) == (_FIRST_IMAGE,) +def test_openai_image_edit_without_a_model_leaves_model_out_of_the_request( + respx_mock: respx.MockRouter, httpx_transport: None +) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( + return_value=httpx.Response(200, json=_EDIT_RESPONSE) + ) + + litellm.image_edit( + prompt="combine the reference images", + image=io.BytesIO(_FIRST_IMAGE), + custom_llm_provider="openai", + api_key="fake-key", + ) + + assert route.call_count == 1 + assert b'name="model"' not in route.calls[0].request.read() + + @pytest.mark.asyncio async def test_openai_image_edit_accepts_bytesio_images(respx_mock: respx.MockRouter, httpx_transport: None) -> None: route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( diff --git a/tests/unit/images/test_main.py b/tests/unit/images/test_main.py index 88d5ae1fc2e..61b2f28b727 100644 --- a/tests/unit/images/test_main.py +++ b/tests/unit/images/test_main.py @@ -300,7 +300,7 @@ async def test_aimage_generation_ignores_a_non_string_base_url_like_the_sync_pat ) response: Final = await litellm.aimage_generation( - model="dall-e-3", + model="gpt-image-2.5-sunburst", prompt="a red circle", api_base=api_base, base_url=httpx.URL(api_base), @@ -357,6 +357,24 @@ def test_image_edit_prices_a_vertex_deployment_at_its_configured_location( assert cost_at("us-central1") == pytest.approx(0.044) +def test_image_edit_accepts_a_string_additional_drop_params(respx_mock: respx.MockRouter) -> None: + edits_route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock( + return_value=httpx.Response(status_code=200, json={"created": 1792551838, "data": [{"b64_json": "aW1n"}]}) + ) + + response: Final = litellm.image_edit( + model="openai/gpt-image-2.5-sunburst", + image=b"\x89PNG\r\n\x1a\nfakepng", + prompt="make the circle blue", + api_key="sk-test", + quality="high", + additional_drop_params="quality", + ) + + assert edits_route.called + assert response.data[0].b64_json == "aW1n" + + class TestCustomLogger(CustomLogger): __test__ = False @@ -670,3 +688,24 @@ async def test_chatgpt_images_forward_client_login_and_workspace( assert "cookie" not in request.headers assert "x-litellm-api-key" not in request.headers assert json.loads(request.content) == {"model": "gpt-image-2", "prompt": "A blue circle"} + + +def test_image_generation_without_a_model_or_provider_fails_before_calling_any_provider( + respx_mock: respx.MockRouter, +) -> None: + with pytest.raises(litellm.BadRequestError, match="model is required"): + litellm.image_generation(prompt="a red circle", api_key="fake_openai_api_key") + + assert not respx_mock.calls + + +def test_topaz_image_variation_works_without_a_model(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post("https://api.topazlabs.com/image/v1/enhance").mock( + return_value=httpx.Response(200, content=b"enhanced-image", headers={"content-type": "image/png"}) + ) + + response: Final = litellm.image_variation(image=b"source-image", custom_llm_provider="topaz", api_key="fake-key") + + assert route.call_count == 1 + assert response.data is not None + assert response.data[0].b64_json == base64.b64encode(b"enhanced-image").decode() diff --git a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py index fdc3771e903..414bcc11602 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_short_circuit.py @@ -6,10 +6,13 @@ and executes the search directly without routing through the backend LLM. """ import json +from typing import Final from unittest.mock import AsyncMock, patch import pytest +import respx +import litellm from litellm.integrations.websearch_interception.handler import ( WebSearchInterceptionLogger, ) @@ -169,6 +172,27 @@ class TestTryShortCircuitSearch: assert result is None + @pytest.mark.asyncio + @pytest.mark.parametrize( + "empty_prompt", + [ + " ", + "Perform a web search for the query: ", + ], + ) + async def test_does_not_short_circuit_empty_or_whitespace_user_message(self, empty_prompt: str): + """User message that is whitespace or empty query → returns None without executing search.""" + logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"]) + + result = await logger.try_short_circuit_search( + model="github_copilot/claude-sonnet-4", + messages=[{"role": "user", "content": empty_prompt}], + tools=[{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}], + custom_llm_provider="github_copilot", + ) + + assert result is None + @pytest.mark.asyncio async def test_search_failure_returns_error_text(self): """Search failure → response with error message, not exception""" @@ -226,6 +250,46 @@ class TestTryShortCircuitSearch: # Query extraction tests # --------------------------------------------------------------------------- +PERPLEXITY_SEARCH_URL: Final = "https://api.perplexity.ai/search" + + +@pytest.fixture +def perplexity_search(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> respx.Route: + monkeypatch.setenv("PERPLEXITYAI_API_KEY", "test-perplexity-key") + monkeypatch.delenv("PERPLEXITY_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + return respx_mock.post(PERPLEXITY_SEARCH_URL).respond( + json={"results": [{"title": "Releases", "url": "https://example.com/releases", "snippet": "Release notes"}]} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("user_message", "searched_query"), + [ + ("Perform a web search for the query: litellm latest release", "litellm latest release"), + ('Perform a web search for the query: "litellm proxy"', '"litellm proxy"'), + ], +) +async def test_short_circuit_searches_the_query_inside_claude_codes_prompt( + perplexity_search: respx.Route, user_message: str, searched_query: str +): + logger: Final = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + + result: Final = await logger.try_short_circuit_search( + model="bedrock/converse/us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": user_message}], + tools=[{"type": "web_search_20250305", "name": "web_search", "max_uses": 8}], + custom_llm_provider="bedrock", + ) + + assert perplexity_search.call_count == 1 + assert json.loads(perplexity_search.calls.last.request.content)["query"] == searched_query + assert result is not None + server_tool_use: Final = next(block for block in result["content"] if block["type"] == "server_tool_use") + assert server_tool_use["input"]["query"] == searched_query + # --------------------------------------------------------------------------- # Integration with entry point diff --git a/tests/unit/interactions/test_agents_http_handler.py b/tests/unit/interactions/test_agents_http_handler.py index 31b78d7a360..77e23e85e47 100644 --- a/tests/unit/interactions/test_agents_http_handler.py +++ b/tests/unit/interactions/test_agents_http_handler.py @@ -22,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.agents.transformation import GeminiAgentsConfig from litellm.llms.gemini.common_utils import GeminiError from litellm.types.agents import ( + AgentCreateRequest, AgentCreateResponse, AgentDeleteResult, AgentListResponse, @@ -104,7 +105,7 @@ class TestCreateAgent: result = handler.create_agent( agents_api_config=config, - name="agent-x", + create_request=AgentCreateRequest(id="agent-x"), litellm_params=litellm_params, logging_obj=logging_obj, extra_headers={"X-Test": "1"}, @@ -117,8 +118,7 @@ class TestCreateAgent: client.post.assert_called_once() kwargs = client.post.call_args.kwargs assert kwargs["url"].endswith("/v1beta/agents") - assert kwargs["json"]["name"] == "agent-x" - assert kwargs["json"]["foo"] == "bar" + assert kwargs["json"] == {"id": "agent-x", "foo": "bar"} assert kwargs["headers"]["X-Test"] == "1" logging_obj.pre_call.assert_called_once() logging_obj.post_call.assert_called_once() @@ -130,7 +130,7 @@ class TestCreateAgent: result = handler.create_agent( agents_api_config=config, - name="agent-x", + create_request=AgentCreateRequest(id="agent-x"), litellm_params=litellm_params, logging_obj=_make_logging_obj(), client=client, @@ -152,7 +152,7 @@ class TestCreateAgent: with pytest.raises(GeminiError): handler.create_agent( agents_api_config=config, - name="agent-x", + create_request=AgentCreateRequest(id="agent-x"), litellm_params=litellm_params, logging_obj=_make_logging_obj(), client=client, @@ -169,7 +169,7 @@ class TestCreateAgent: result = await handler.async_create_agent( agents_api_config=config, - name="agent-y", + create_request=AgentCreateRequest(id="agent-y"), litellm_params=litellm_params, logging_obj=_make_logging_obj(), extra_body={"baz": "qux"}, @@ -193,7 +193,7 @@ class TestCreateAgent: with pytest.raises(GeminiError): await handler.async_create_agent( agents_api_config=config, - name="agent-y", + create_request=AgentCreateRequest(id="agent-y"), litellm_params=litellm_params, logging_obj=_make_logging_obj(), client=client, diff --git a/tests/unit/interactions/test_agents_main_and_utils.py b/tests/unit/interactions/test_agents_main_and_utils.py index 395801ff059..97d93563699 100644 --- a/tests/unit/interactions/test_agents_main_and_utils.py +++ b/tests/unit/interactions/test_agents_main_and_utils.py @@ -9,9 +9,14 @@ small helper utilities without touching the network. """ import asyncio +import json +from collections.abc import Callable, Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm @@ -137,7 +142,7 @@ class TestSyncEntryPoints: assert response is sentinel handler.create_agent.assert_called_once() kw = handler.create_agent.call_args.kwargs - assert kw["name"] == "waverunner" + assert kw["create_request"].id == "waverunner" assert kw["_is_async"] is False assert kw["extra_headers"] == {"X-Test": "1"} assert kw["extra_body"] == {"foo": "bar"} @@ -205,7 +210,7 @@ class TestAsyncEntryPoints: def fake_create_agent(**kwargs): assert kwargs["_is_async"] is True - assert kwargs["name"] == "waverunner" + assert kwargs["create_request"].id == "waverunner" return sentinel handler = MagicMock() @@ -349,3 +354,250 @@ class TestAsyncErrorWrapping: with patch(_HANDLER_PATH, handler): with pytest.raises(litellm.APIConnectionError): await alist_versions(name="waverunner", api_key="AIza") + + +_GEMINI_AGENTS_URL: Final = "https://generativelanguage.googleapis.com/v1beta/agents" + + +def _sent_json(route: respx.Route) -> object: + return json.loads(route.calls.last.request.content) + + +class TestCreateRequestBody: + def test_create_sends_every_current_field_verbatim(self, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "sentinel", "base_agent": "antigravity-preview-09-2026"}) + ) + + response: Final = create( + id="sentinel", + base_agent="antigravity-preview-09-2026", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash", "max_total_tokens": None}, + system_instruction="Always answer in exactly three words.", + base_environment={"type": "remote", "sources": [{"type": "gcs", "source": "gs://bucket/skill"}]}, + description="Watches the queue.", + tools=[{"type": "code_execution"}], + api_key="AIza-test", + ) + + assert _sent_json(route) == { + "id": "sentinel", + "base_agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash", "max_total_tokens": None}, + "system_instruction": "Always answer in exactly three words.", + "base_environment": {"type": "remote", "sources": [{"type": "gcs", "source": "gs://bucket/skill"}]}, + "description": "Watches the queue.", + "tools": [{"type": "code_execution"}], + } + assert (response.id, response.name) == ("sentinel", "sentinel") + + @pytest.mark.parametrize( + ("identifier", "instruction", "expected"), + [ + pytest.param( + {"name": "slides"}, + {"instructions": "Make slides."}, + {"id": "slides", "system_instruction": "Make slides."}, + id="preview-aliases", + ), + pytest.param( + {"id": "slides", "name": "slides"}, + {"system_instruction": "Make slides.", "instructions": "Make slides."}, + {"id": "slides", "system_instruction": "Make slides."}, + id="same-value-under-both-names", + ), + ], + ) + def test_create_sends_preview_aliases_under_current_names( + self, + respx_mock: respx.MockRouter, + identifier: Mapping[str, str], + instruction: Mapping[str, str], + expected: Mapping[str, str], + ) -> None: + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "slides"}) + ) + + create(**identifier, **instruction, base_agent="antigravity-preview-09-2026", api_key="AIza-test") + + assert _sent_json(route) == {**expected, "base_agent": "antigravity-preview-09-2026"} + + @pytest.mark.parametrize( + "fields", + [ + pytest.param({"base_agent": "antigravity-preview-09-2026"}, id="no-identifier"), + pytest.param({"id": "slides", "name": "decks"}, id="different-id-and-name"), + pytest.param( + {"id": "slides", "system_instruction": "Make slides.", "instructions": "Make decks."}, + id="different-system-instruction-and-instructions", + ), + pytest.param({"id": "slides", "agent_config": "antigravity"}, id="agent-config-not-an-object"), + ], + ) + def test_create_rejects_a_malformed_identifier_or_instruction( + self, respx_mock: respx.MockRouter, fields: Mapping[str, str] + ) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + create(**fields, api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert not respx_mock.calls + + def test_create_with_drop_params_drops_only_the_invalid_optional_fields( + self, respx_mock: respx.MockRouter + ) -> None: + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "slides"}) + ) + + create( + id="slides", + agent_config="antigravity", + tools="code_execution", + instructions="Make slides.", + drop_params=True, + api_key="AIza-test", + ) + + assert _sent_json(route) == {"id": "slides", "system_instruction": "Make slides."} + + def test_create_with_drop_params_drops_a_clashing_instructions_alias(self, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "slides"}) + ) + + create( + id="slides", + system_instruction="Make slides.", + instructions="Make decks.", + drop_params=True, + api_key="AIza-test", + ) + + assert _sent_json(route) == {"id": "slides", "system_instruction": "Make slides."} + + def test_create_with_drop_params_still_rejects_a_clashing_identifier(self, respx_mock: respx.MockRouter) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + create(id="slides", name="decks", drop_params=True, api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert not respx_mock.calls + + @pytest.mark.parametrize( + ("fields", "expected_message"), + [ + pytest.param( + {"id": "slides", "instructions": 5, "tools": "code_execution"}, + "Invalid agent create request: 'instructions' must be a string; 'tools' must be a list of objects." + " Set drop_params to drop invalid optional fields.", + id="names-the-alias-the-caller-sent", + ), + pytest.param( + {"name": 7, "base_agent": "antigravity-preview-09-2026"}, + "Invalid agent create request: 'name' must be a string.", + id="names-the-id-alias-the-caller-sent", + ), + pytest.param( + {"id": 7, "agent_config": "antigravity", "drop_params": True}, + "Invalid agent create request: 'agent_config' must be an object; 'id' must be a string.", + id="invalid-id-is-never-dropped", + ), + ], + ) + def test_create_rejects_invalid_fields_by_the_names_the_caller_sent( + self, respx_mock: respx.MockRouter, fields: Mapping[str, object], expected_message: str + ) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + create(**fields, api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert exc_info.value.message.endswith(expected_message) + assert not respx_mock.calls + + @pytest.mark.asyncio + async def test_acreate_sends_every_current_field_verbatim( + self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + route: Final = respx_mock.post(_GEMINI_AGENTS_URL).mock( + return_value=httpx.Response(200, json={"id": "sentinel"}) + ) + + response: Final = await acreate( + id="sentinel", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, + system_instruction="Always answer in exactly three words.", + api_key="AIza-test", + ) + + assert _sent_json(route) == { + "id": "sentinel", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + "system_instruction": "Always answer in exactly three words.", + } + assert response.id == "sentinel" + + +class TestAgentIdentifier: + def test_get_by_id_fetches_that_agent(self, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.get(f"{_GEMINI_AGENTS_URL}/slides").mock( + return_value=httpx.Response(200, json={"id": "slides", "base_agent": "antigravity-preview-09-2026"}) + ) + + response: Final = get(id="slides", api_key="AIza-test") + + assert route.call_count == 1 + assert (response.id, response.base_agent) == ("slides", "antigravity-preview-09-2026") + + def test_delete_by_id_deletes_that_agent(self, respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.delete(f"{_GEMINI_AGENTS_URL}/slides").mock( + return_value=httpx.Response(200, json={}) + ) + + result: Final = delete(id="slides", api_key="AIza-test") + + assert route.call_count == 1 + assert (result.name, result.deleted) == ("slides", True) + + def test_list_versions_by_id_lists_that_agents_versions(self, respx_mock: respx.MockRouter) -> None: + version: Final = {"agent": "slides", "name": "agents/slides/versions/v1"} + route: Final = respx_mock.get(f"{_GEMINI_AGENTS_URL}/slides/versions").mock( + return_value=httpx.Response(200, json={"agentVersions": [version]}) + ) + + response: Final = list_versions(id="slides", api_key="AIza-test") + + assert route.call_count == 1 + assert response.agent_versions == [version] + + @pytest.mark.parametrize( + "operation", + [pytest.param(get, id="get"), pytest.param(delete, id="delete"), pytest.param(list_versions, id="list_versions")], + ) + @pytest.mark.parametrize( + "identifier", + [ + pytest.param({}, id="no-identifier"), + pytest.param({"id": "slides", "name": "decks"}, id="different-id-and-name"), + ], + ) + def test_rejects_a_missing_or_conflicting_identifier( + self, + respx_mock: respx.MockRouter, + operation: Callable[..., object], + identifier: Mapping[str, str], + ) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + operation(**identifier, api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert not respx_mock.calls + + @pytest.mark.asyncio + async def test_aget_rejects_a_missing_identifier(self, respx_mock: respx.MockRouter) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + await aget(api_key="AIza-test") + + assert exc_info.value.status_code == 400 + assert not respx_mock.calls diff --git a/tests/unit/interactions/test_interactions_streaming_iterator.py b/tests/unit/interactions/test_interactions_streaming_iterator.py index afcd0d559d7..fb1fe92d637 100644 --- a/tests/unit/interactions/test_interactions_streaming_iterator.py +++ b/tests/unit/interactions/test_interactions_streaming_iterator.py @@ -16,7 +16,10 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.interactions import streaming_iterator as interactions_streaming_iterator_module -from litellm.interactions.streaming_iterator import InteractionsAPIStreamingIterator +from litellm.interactions.streaming_iterator import ( + InteractionsAPIStreamingIterator, + SyncInteractionsAPIStreamingIterator, +) from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig @@ -148,3 +151,43 @@ async def test_stream_logs_its_first_event_time_measured_from_the_call_start(): assert first_event_at is not None assert recorder.logged == [(call_start, first_event_at)] + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _first_event( + iterator: InteractionsAPIStreamingIterator | SyncInteractionsAPIStreamingIterator, +) -> InteractionsAPIStreamingResponse: + if isinstance(iterator, InteractionsAPIStreamingIterator): + return await iterator.__anext__() + return next(iterator) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("iterator_cls", [InteractionsAPIStreamingIterator, SyncInteractionsAPIStreamingIterator]) +async def test_stream_keeps_unicode_line_separators_inside_the_delta_text( + iterator_cls: type[InteractionsAPIStreamingIterator] | type[SyncInteractionsAPIStreamingIterator], +) -> None: + delta: Final = {"type": "text", "text": _LINE_SEPARATED_TEXT} + event: Final = {"event_type": "step.delta", "id": "interaction-1", "delta": delta} + iterator: Final = iterator_cls( + response=httpx.Response( + 200, stream=httpx.ByteStream(f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode()) + ), + model="gemini/gemini-3.1-pro-preview", + interactions_api_config=GoogleAIStudioInteractionsConfig(), + logging_obj=LitellmLogging( + model="gemini/gemini-3.1-pro-preview", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="ainteraction", + start_time=datetime(2026, 10, 9, 12, 0, 0), + litellm_call_id="interactions-separators", + function_id="interactions-separators", + ), + ) + + first: Final = await _first_event(iterator) + + assert first.delta == delta diff --git a/tests/unit/interactions/test_main.py b/tests/unit/interactions/test_main.py index 123f44ca779..5af5dd7d82a 100644 --- a/tests/unit/interactions/test_main.py +++ b/tests/unit/interactions/test_main.py @@ -236,6 +236,159 @@ class TestGoogleInteractionsCreate: assert json.loads(route.calls.last.request.content) == expected_body assert response.status == "completed" + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_with_agent_config_sends_it_to_gemini(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response( + 200, + json={"id": "interaction-agent", "object": "interaction", "status": "completed"}, + ) + ) + + interactions.create( + agent="antigravity-preview-09-2026", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, + input="Say hi in one word.", + environment="remote", + api_key="gemini-offline", + ) + + assert json.loads(route.calls.last.request.content) == { + "agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + "input": "Say hi in one word.", + "environment": "remote", + } + + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_keeps_positional_response_settings_in_place(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response(200, json={"id": "interaction-model", "status": "completed"}) + ) + + interactions.create( + "gemini/gemini-2.5-flash", + None, + "Describe a cat.", + None, + None, + None, + None, + None, + None, + None, + ["TEXT"], + api_key="gemini-offline", + ) + + assert json.loads(route.calls.last.request.content) == { + "model": "gemini-2.5-flash", + "input": "Describe a cat.", + "response_modalities": ["TEXT"], + } + + @pytest.mark.parametrize( + ("target", "agent_config", "expected_target"), + [ + pytest.param( + {"model": "gemini/gemini-2.5-flash"}, + {"type": "antigravity", "model": "gemini-3.8-flash"}, + {"model": "gemini-2.5-flash"}, + id="model-drops-it", + ), + pytest.param( + {"agent": "antigravity-preview-09-2026"}, + {"type": "antigravity", "model": "gemini-3.8-flash"}, + { + "agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + }, + id="agent-keeps-an-object", + ), + pytest.param( + {"agent": "antigravity-preview-09-2026"}, + "antigravity", + {"agent": "antigravity-preview-09-2026"}, + id="agent-drops-a-non-object", + ), + pytest.param( + {"agent": "antigravity-preview-09-2026"}, + {"model": "gemini-3.8-flash"}, + {"agent": "antigravity-preview-09-2026"}, + id="agent-drops-an-object-without-a-type", + ), + pytest.param( + {"agent": "antigravity-preview-09-2026"}, + {"type": 7, "model": "gemini-3.8-flash"}, + {"agent": "antigravity-preview-09-2026"}, + id="agent-drops-a-non-string-type", + ), + ], + ) + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_with_drop_params_drops_agent_config_gemini_cannot_take( + self, + respx_mock: MockRouter, + target: Mapping[str, str], + agent_config: object, + expected_target: Mapping[str, object], + ) -> None: + route: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response(200, json={"id": "interaction-drop", "status": "completed"}) + ) + + interactions.create( + **target, + agent_config=agent_config, + input="Say hi in one word.", + drop_params=True, + api_key="gemini-offline", + ) + + assert json.loads(route.calls.last.request.content) == {**expected_target, "input": "Say hi in one word."} + + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_rejects_agent_config_a_bridged_provider_cannot_take(self, respx_mock: MockRouter) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_openai_response_body("resp-agent-config", "Hi")) + ) + + with pytest.raises(litellm.UnsupportedParamsError, match="'agent_config' is not supported for openai") as e: + interactions.create( + model="gpt-5.6", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, + input="Say hi in one word.", + api_key="sk-offline", + ) + + assert e.value.status_code == 400 + assert not route.called + + @pytest.mark.parametrize( + "drop_kwargs", + [ + pytest.param({"drop_params": True}, id="drop-params"), + pytest.param({"additional_drop_params": ["agent_config"]}, id="additional-drop-params"), + ], + ) + @pytest.mark.usefixtures("fake_provider_credentials") + def test_create_with_drop_params_drops_agent_config_a_bridged_provider_cannot_take( + self, respx_mock: MockRouter, drop_kwargs: Mapping[str, object] + ) -> None: + route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_openai_response_body("resp-agent-config", "Hi")) + ) + + interactions.create( + model="gpt-5.6", + agent_config={"type": "antigravity", "model": "gemini-3.8-flash"}, + input="Say hi in one word.", + api_key="sk-offline", + **drop_kwargs, + ) + + assert json.loads(route.calls.last.request.content) == {"model": "gpt-5.6", "input": "Say hi in one word."} + @pytest.mark.parametrize( ("model", "api_key", "url", "expected_body", "sse_body", "expected_chunks"), [ diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index d50ab908b43..b9675f6fd5c 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( schema_has_unresolvable_ref, schema_property_count, split_concatenated_json_objects, + split_media_for_tpm, strip_encrypted_reasoning_from_messages, strip_litellm_response_only_fields, system_messages_first, @@ -2512,3 +2513,118 @@ def test_shared_image_url_agrees_with_responses_normalization( expected: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE"}} assert as_openai_image_part(part) == expected assert as_openai_image_part(normalized) == expected + + +@pytest.mark.parametrize( + "part,projected,media", + [ + ( + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}}, + {"type": "image_url", "image_url": {"url": "", "detail": "high"}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}}, + ), + ( + {"type": "input_image", "image_url": "https://example.com/image.png", "detail": "low"}, + {"type": "input_image", "image_url": "", "detail": "low"}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "low"}}, + ), + ( + {"type": "computer_screenshot", "image_url": "data:image/png;base64,IMAGE"}, + {"type": "computer_screenshot", "image_url": ""}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}}, + ), + ( + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "IMAGE"}}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": ""}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}}, + ), + ( + {"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}}, + {"type": "input_audio", "input_audio": {"data": "", "format": "wav"}}, + {"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}}, + ), + ( + {"type": "file", "file": {"file_data": "IMAGE", "format": "image/png", "filename": "image.png"}}, + {"type": "file", "file": {"file_data": "", "format": "image/png", "filename": "image.png"}}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}}, + ), + ( + {"type": "input_file", "file_data": "data:audio/wav;base64,AUDIO", "filename": "audio.wav"}, + {"type": "input_file", "file_data": "", "filename": "audio.wav"}, + {"type": "input_audio", "input_audio": {"data": "AUDIO", "format": "wav"}}, + ), + ], +) +def test_tpm_projection_blanks_only_declared_media_transport( + part: Mapping[str, object], projected: Mapping[str, object], media: Mapping[str, object], +) -> None: + text: Final = {"type": "text", "text": "data:image/png;base64,IMAGE"} + request: Final = [{"role": "user", "content": [text, {**part, "cache_control": {"type": "ephemeral"}}]}] + original: Final = copy.deepcopy(request) + assert split_media_for_tpm(request) == ( + [{"role": "user", "content": [text, {**projected, "cache_control": {"type": "ephemeral"}}]}], + (media,), + ) + assert request == original + + +def test_tpm_projection_retains_nested_and_native_history_around_media() -> None: + image: Final = {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "high"}} + empty_image: Final = {"type": "image_url", "image_url": {"url": "", "detail": "high"}} + document: Final = {"type": "document", "title": "keep", "source": {"type": "content", "content": [image]}} + empty_document: Final = {**document, "source": {"type": "content", "content": [empty_image]}} + tool: Final = {"type": "tool_result", "tool_use_id": "call", "content": [document, {"type": "text", "text": "keep"}]} + empty_tool: Final = {**tool, "content": [empty_document, {"type": "text", "text": "keep"}]} + request: Final = [ + {"role": "user", "content": [tool]}, + {"type": "function_call_output", "call_id": "call", "output": [image]}, + {"type": "custom_tool_call_output", "call_id": "custom", "output": [image]}, + {"type": "computer_call_output", "call_id": "computer", "output": { + "type": "computer_screenshot", "image_url": "data:image/png;base64,IMAGE", + }}, + {"type": "image_generation_call", "id": "ig", "result": "IMAGE", "status": "completed"}, + {"type": "code_interpreter_call", "id": "ci", "code": "keep()", "outputs": [ + {"type": "logs", "logs": "keep"}, {"type": "image", "url": "https://example.com/image.png"}, + ]}, + ] + original: Final = copy.deepcopy(request) + projected, media = split_media_for_tpm(request) + assert projected == [ + {"role": "user", "content": [empty_tool]}, + {"type": "function_call_output", "call_id": "call", "output": [empty_image]}, + {"type": "custom_tool_call_output", "call_id": "custom", "output": [empty_image]}, + {"type": "computer_call_output", "call_id": "computer", "output": { + "type": "computer_screenshot", "image_url": "", + }}, + {"type": "image_generation_call", "id": "ig", "result": "", "status": "completed"}, + {"type": "code_interpreter_call", "id": "ci", "code": "keep()", "outputs": [ + {"type": "logs", "logs": "keep"}, {"type": "image", "url": ""}, + ]}, + ] + assert media == ( + image, image, image, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,IMAGE", "detail": "auto"}}, + {"type": "image_url", "image_url": {"url": "data:application/octet-stream;base64,IMAGE", "detail": "auto"}}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.png", "detail": "auto"}}, + ) + assert request == original + + +def test_tpm_projection_preserves_non_media_identity_and_opaque_values() -> None: + image: Final = {"type": "input_image", "image_url": "data:image/png;base64,OPAQUE"} + request: Final = [ + {"role": "user", "content": json.dumps(image)}, + {"type": "custom_tool_call_output", "call_id": "call", "output": json.dumps(image)}, + {"type": "function_call", "call_id": "call", "name": "keep", "arguments": {"content": [image]}}, + {"type": "file_search_call", "results": [{"text": "keep", "content": [image]}]}, + {"type": "image_generation_call", "id": "ig", "status": "in_progress", "result": None}, + {"role": "user", "content": [ + {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "PDF"}}, + {"type": "document", "source": {"type": "text", "data": json.dumps(image)}}, + {"type": "file", "file": {"format": "application/pdf", "file_data": "PDF"}}, + {"type": "text", "text": "keep", "metadata": {"content": [image]}}, + ]}, + ] + projected, media = split_media_for_tpm(request) + assert projected is request + assert media == () diff --git a/tests/unit/litellm_core_utils/test_audio_utils.py b/tests/unit/litellm_core_utils/test_audio_utils.py index 155f6680416..122445e8e30 100644 --- a/tests/unit/litellm_core_utils/test_audio_utils.py +++ b/tests/unit/litellm_core_utils/test_audio_utils.py @@ -2,16 +2,23 @@ Test the audio utils functionality in litellm_core_utils/audio_utils/utils.py """ +import base64 import io import os import tempfile +from typing import Final from unittest.mock import mock_open, patch import pytest +import soundfile as sf from litellm.litellm_core_utils.audio_utils.utils import ( + AUDIO_BYTES_PER_TOKEN, + AUDIO_TOKENS_PER_SECOND, + DEFAULT_AUDIO_TOKEN_ESTIMATE, ProcessedAudioFile, calculate_request_duration, + estimate_audio_tokens, get_audio_file_content_hash, get_audio_file_for_health_check, get_audio_file_name, @@ -19,6 +26,63 @@ from litellm.litellm_core_utils.audio_utils.utils import ( ) +def _encoded_audio(seconds: int, sample_rate: int, channels: int, container: str) -> str: + buffer: Final = io.BytesIO() + with sf.SoundFile(buffer, "w", samplerate=sample_rate, channels=channels, format=container) as audio: + audio.buffer_write(b"\0" * (seconds * sample_rate * channels * 2), dtype="int16") + return base64.b64encode(buffer.getvalue()).decode() + + +@pytest.mark.parametrize( + ("sample_rate", "channels", "container"), [(8000, 1, "WAV"), (48000, 2, "WAV"), (24000, 1, "FLAC")] +) +def test_audio_estimate_uses_duration_independent_of_sample_rate_channels_and_compression( + sample_rate: int, channels: int, container: str +) -> None: + assert estimate_audio_tokens(_encoded_audio(1, sample_rate, channels, container)) == AUDIO_TOKENS_PER_SECOND + assert estimate_audio_tokens(_encoded_audio(31, sample_rate, channels, container)) == 31 * AUDIO_TOKENS_PER_SECOND + + +# Audio rates verified 2026-10-28: https://ai.google.dev/gemini-api/docs/generate-content/media-resolution and +# https://developers.openai.com/api/docs/guides/voice-latency-cost +@pytest.mark.parametrize( + ("models", "fallback", "expected"), + [ + (("openai/gpt-audio-1.5",), "gemini-3.8-flash", 10), + (("gpt-audio-1.5", "gemini/gemini-3.8-flash"), None, 25), + (("vertex_ai/gemini-3.8-flash",), None, 25), + (("openrouter/google/gemini-3.8-flash",), None, 25), + (("azure/audio-deployment",), "gpt-audio-1.5", 10), + (("opaque-deployment",), "gemini-3.8-flash", 25), + ((), "gemini-3.8-flash", 25), + (("gemini/gemini-2.5-flash", "gemini/gemini-3-pro-preview"), None, 32), + (("other/gemini-3.8-flash",), None, 10), + ], +) +def test_audio_estimate_prefers_each_backend_rate_over_the_requested_model( + models: tuple[str, ...], fallback: str | None, expected: int +) -> None: + assert estimate_audio_tokens(_encoded_audio(1, 8000, 1, "WAV"), models, fallback) == expected + + +@pytest.mark.parametrize("data", [None, "", "invalid-base64!", base64.b64encode(b"unreadable audio").decode()]) +def test_audio_estimate_retains_nonzero_fallback_for_unreadable_audio(data: str | None) -> None: + assert estimate_audio_tokens(data) == DEFAULT_AUDIO_TOKEN_ESTIMATE + + +def test_audio_estimate_retains_size_based_fallback_for_large_unreadable_audio() -> None: + data: Final = base64.b64encode(b"\0" * (AUDIO_BYTES_PER_TOKEN * 400)).decode() + assert estimate_audio_tokens(data) == 400 + + +def test_audio_estimate_bounds_base64_allocation_before_decoding() -> None: + data: Final = base64.b64encode(b"\0" * (16 * 1024 * 1024 + 1)).decode() + with patch("base64.b64decode", side_effect=AssertionError("oversized audio must not be decoded")) as decode: + estimate: Final = estimate_audio_tokens(data) + assert estimate == len(data) * 3 // 4 // AUDIO_BYTES_PER_TOKEN + decode.assert_not_called() + + class TestProcessAudioFile: """Test the process_audio_file function with various input types""" diff --git a/tests/unit/router_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py similarity index 97% rename from tests/unit/router_utils/test_batch_utils.py rename to tests/unit/litellm_core_utils/test_batch_utils.py index 79fe86a37c0..cd37d635c43 100644 --- a/tests/unit/router_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -7,7 +7,7 @@ from typing import Dict, Final, List import pytest import litellm -from litellm.router_utils.batch_utils import ( +from litellm.litellm_core_utils.batch_utils import ( InMemoryFile, get_router_metadata_variable_name, parse_jsonl_with_embedded_newlines, @@ -116,7 +116,7 @@ def test_router_metadata_variable_name(): def test_non_json_input(): """Test that replace_model_in_jsonl returns original content for non-JSON input""" - from litellm.router_utils.batch_utils import replace_model_in_jsonl + from litellm.litellm_core_utils.batch_utils import replace_model_in_jsonl # Test with non-JSON string non_json_str = "This is not a JSON string" @@ -138,7 +138,7 @@ def test_non_json_input(): def test_should_replace_model_in_jsonl(): """Test that should_replace_model_in_jsonl returns the correct value""" - from litellm.router_utils.batch_utils import should_replace_model_in_jsonl + from litellm.litellm_core_utils.batch_utils import should_replace_model_in_jsonl assert should_replace_model_in_jsonl(purpose="batch") is True assert should_replace_model_in_jsonl(purpose="batch", passthrough=True) is False @@ -340,7 +340,7 @@ def test_replace_model_in_jsonl_with_embedded_newlines(): def test_is_batch_retrieve_call_type_matches_only_batch_retrieves(): - from litellm.router_utils.batch_utils import is_batch_retrieve_call_type + from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import CallTypes assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True @@ -389,7 +389,8 @@ def setup_and_teardown(): ) def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None: from openai.types import Batch - from litellm.router_utils.batch_utils import batch_id_to_charge_once + + from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once batch: Final = Batch( id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", @@ -401,7 +402,7 @@ def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, def test_batch_charge_markers_keep_counters_and_batches_independent() -> None: from litellm.caching.in_memory_cache import InMemoryCache - from litellm.router_utils.batch_utils import batch_charged_once_marker_key + from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key cache: Final = InMemoryCache() cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1) diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index d66c37ba8c2..66a8e6bbda6 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -256,6 +256,7 @@ class TestMapFinishReasonAnthropic: ("max_tokens", "length"), ("tool_use", "tool_calls"), ("compaction", "length"), + ("model_context_window_exceeded", "length"), ("content_filtered", "content_filter"), ], ) diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 9e8e9cf2a80..55e7c47e2c9 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -1516,9 +1516,9 @@ def test_bedrock_mantle_openai_envelope_context_overflow_maps_to_context_window_ def test_bedrock_mantle_streamed_context_overflow_event_maps_to_context_window_exceeded(): - from litellm.responses.streaming_iterator import _map_stream_error_to_exception + from litellm.responses.streaming_iterator import map_stream_error_to_exception - mapped_exception = _map_stream_error_to_exception( + mapped_exception = map_stream_error_to_exception( { "code": "context_length_exceeded", "message": "Your input exceeds the context window of this model. Please adjust your input and try again.", diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index b584a72855a..88e2bc4b1fd 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -723,6 +723,37 @@ async def test_ahealth_check_probes_hosted_vllm_with_a_choice_question( } +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +async def test_ahealth_check_probes_ollama_decisions_with_one_noul_question_without_auth( + provider: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond( + json={ + "model": "d3:4b", + "answers": {"reachable": {"type": "noul", "noul": 1.0}}, + "usage": {"input_tokens": 12, "output_tokens": 1}, + } + ) + + result: Final = await ahealth_check( + {"model": f"{provider}/d3:4b", "api_base": "http://ollama.local:11434"}, mode="decisions" + ) + + assert "error" not in result, result + assert upstream.call_count == 1 + assert "authorization" not in upstream.calls.last.request.headers + assert json.loads(upstream.calls.last.request.content) == { + "model": "d3:4b", + "state": "health check", + "questions": {"reachable": {"type": "noul", "instructions": "Is the service reachable?"}}, + } + + @pytest.mark.asyncio @pytest.mark.parametrize("mode", ("evaluation", "decisions")) async def test_ahealth_check_keeps_the_noul_probe_for_other_decisions_providers( diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index e786753e1a8..299088a6c4f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -13775,3 +13775,35 @@ def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_fro logging_obj.record_native_bedrock_guardrail({"trace": trace}) assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"] + + +def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_calls_share_one_metadata_dict() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost + + shared_metadata = {"tag": "x"} + traces = tuple( + {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": units}}}}}} + for units in (1, 2) + ) + costs = [] + for call_index, trace in enumerate(traces): + logging_obj = Logging( + model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id=f"call-shared-metadata-{call_index}", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": shared_metadata} + logging_obj.record_native_bedrock_guardrail({"trace": trace}) + recorded = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] + costs.append(guardrail_information_cost(recorded)) + + (first_cost, second_cost) = costs + assert second_cost == pytest.approx(2 * first_cost) + assert shared_metadata == {"tag": "x"} diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 6d700bae514..79978ecfe8e 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1,5 +1,6 @@ import asyncio import datetime +import itertools import json import queue import sys @@ -13,6 +14,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest +from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, TypeAdapter import litellm @@ -28,6 +30,7 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkParsed, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -503,6 +506,118 @@ def test_multi_chunk_reasoning_and_content( assert initialized_custom_stream_wrapper.sent_last_thinking_block is True +def _stream_chunk(delta: dict[str, object], finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream(choices=[{"index": 0, "delta": delta, "finish_reason": finish_reason}]) + + +def _content_through_first_call(deltas: Sequence[Delta]) -> str: + first_call: Final = next( + (index for index, delta in enumerate(deltas) if delta.tool_calls or delta.function_call), len(deltas) + ) + return "".join(delta.content or "" for delta in deltas[: first_call + 1]) + + +_GET_WEATHER_CALL: Final = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.parametrize( + "chunks,expected_content,expected_tool_calls", + [ + ( + ( + _stream_chunk({"reasoning_content": "Check the weather"}), + _stream_chunk({"tool_calls": [_GET_WEATHER_CALL]}), + _stream_chunk({}, "tool_calls"), + ), + "Check the weather", + 1, + ), + ( + ( + _stream_chunk({"reasoning_content": "Pick a function"}), + _stream_chunk({"function_call": {"name": "get_weather", "arguments": "{}"}}), + _stream_chunk({}, "function_call"), + ), + "Pick a function", + 0, + ), + ( + ( + _stream_chunk({"reasoning_content": "Step one"}), + _stream_chunk({"reasoning_content": ", step two"}), + _stream_chunk({}, "length"), + ), + "Step one, step two", + 0, + ), + ( + ( + _stream_chunk({"reasoning_content": "a"}), + _stream_chunk({"content": "b"}), + _stream_chunk({"reasoning_content": "c"}), + _stream_chunk({"content": "d"}), + _stream_chunk({}, "stop"), + ), + "abcd", + 0, + ), + ( + (_stream_chunk({"reasoning_content": "Still thinking"}),), + "Still thinking", + 0, + ), + ( + (_stream_chunk({"content": "hi"}), _stream_chunk({}, "stop")), + "hi", + 0, + ), + ], + ids=[ + "reasoning_then_tool_call", + "reasoning_then_function_call", + "reasoning_cut_by_length", + "reasoning_after_answer", + "stream_ends_mid_reasoning", + "answer_without_reasoning", + ], +) +async def test_merged_reasoning_stream_keeps_think_tags_balanced( + logging_obj: Logging, + is_async: bool, + chunks: tuple[ModelResponseStream, ...], + expected_content: str, + expected_tool_calls: int, +) -> None: + response: Final = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=[chunk.model_copy(deep=True) for chunk in chunks]), + model="deepseek-reasoner", + logging_obj=logging_obj, + ) + response.merge_reasoning_content_in_choices = True + + deltas: Final = ( + tuple([chunk.choices[0].delta async for chunk in response]) + if is_async + else tuple(chunk.choices[0].delta for chunk in response) + ) + logged: Final = litellm.stream_chunk_builder(chunks=response.chunks) + + assert "".join(delta.content or "" for delta in deltas) == expected_content + assert _content_through_first_call(deltas) == expected_content + assert sum(len(delta.tool_calls or ()) for delta in deltas) == expected_tool_calls + assert logged.choices[0].message.content == (expected_content or None) + assert any(chunk.choices[0].finish_reason for chunk in response.chunks if chunk.choices) == any( + chunk.choices[0].finish_reason for chunk in chunks + ) + + def test_strip_sse_data_from_chunk(): """Test the static method that strips 'data: ' prefix from SSE chunks""" # Test with string inputs @@ -1367,6 +1482,148 @@ async def test_bedrock_validation_error_raises_directly(logging_obj: Logging): assert getattr(excinfo.value, "status_code", None) == 400 +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +@pytest.mark.parametrize("prefix", ["", "Partial answer"]) +@pytest.mark.parametrize("abort_index", [0, 1]) +@pytest.mark.parametrize("abort_content", ["", "unfinished"]) +async def test_aborted_stream_raises_instead_of_finishing_successfully( + async_stream: bool, transport: str, prefix: str, abort_index: int, abort_content: str, logging_obj: Logging +) -> None: + initial: Final = ( + (({"index": 0, "delta": {"role": "assistant", "content": prefix}, "finish_reason": None},),) if prefix else () + ) + final: Final = tuple( + { + "index": index, + "delta": {"content": abort_content}, + "finish_reason": "abort" if index == abort_index else "stop", + } + for index in range(abort_index + 1) + ) + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-aborted", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + "choices": choices, + } + ) + + "\n\n" + for choices in (*initial, final) + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + if transport == "sdk" + else OpenAIChatCompletionStreamingHandler( + streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), + sync_stream=not async_stream, + ) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", + ) + + try: + with pytest.raises(MidStreamFallbackError, match="abort") as excinfo: + _chunks: Final = [chunk async for chunk in response] if async_stream else list(response) + + assert excinfo.value.is_pre_first_chunk is (not prefix) + assert excinfo.value.generated_content == prefix + assert getattr(excinfo.value.original_exception, "status_code", None) == 500 + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +async def test_stream_with_null_choices_chunk_finishes_normally( + async_stream: bool, transport: str, logging_obj: Logging +) -> None: + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "chatcmpl-null-choices", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + "choices": choices, + **extra, + } + ) + + "\n\n" + for choices, extra in ( + ([{"index": 0, "delta": {"role": "assistant", "content": "hi"}, "finish_reason": None}], {}), + ([{"index": 0, "delta": {}, "finish_reason": "stop"}], {}), + (None, {"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}}), + ) + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + if transport == "sdk" + else OpenAIChatCompletionStreamingHandler( + streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), + sync_stream=not async_stream, + ) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", + ) + + try: + chunks: Final = [chunk async for chunk in response] if async_stream else list(response) + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + streamed_choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in streamed_choices) == "hi" + assert [choice.finish_reason for choice in streamed_choices if choice.finish_reason] == ["stop"] + + def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper: """A CustomStreamWrapper over the real OpenAI-compatible line iterator, fed an HTTP 200 SSE body that carries an in-body error payload the way @@ -7704,7 +7961,7 @@ def gemini_mock_post_streaming(url, **kwargs): mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "text/event-stream"} - mock_response.iter_lines = MagicMock(return_value=stream_response()) + mock_response.iter_text = MagicMock(return_value=stream_response()) return mock_response diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 92f971f64e9..b02e0732817 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.token_counter import ( _get_extrapolating_count_function, _get_tiktoken_count_function, calculate_img_tokens, + estimate_image_tokens, high_detail_image_token_upper_bound, offload_token_count, ) @@ -2152,6 +2153,136 @@ def test_high_detail_image_token_upper_bound_is_reached_by_the_largest_high_res_ assert calculate_img_tokens(_png_data_url(1, 1), mode="high") < high_detail_image_token_upper_bound() +# Image sizing verified 2026-10-28: https://developers.openai.com/api/docs/guides/images-vision +# and https://platform.claude.com/docs/en/build-with-claude/vision-coordinates +@pytest.mark.parametrize( + ("model", "width", "height", "detail", "expected"), + [ + ("gpt-4o-mini", 1024, 1024, "low", 2833), + ("gpt-4o-mini", 1024, 1024, "high", 25501), + ("gpt-5.1", 512, 512, "auto", 210), + ("gpt-6-astra", 2048, 2048, "high", 3000), + ("gpt-6-astra", 4096, 512, "high", 2458), + ("gpt-6-astra", 2048, 2048, "auto", 4916), + ("gpt-5.6-sol", 2048, 2048, "low", 308), + ("gpt-5.5", 6400, 6400, "original", 12000), + ("gpt-5.4-mini", 2048, 2048, "low", 4916), + ("gpt-5.4-nano", 2048, 2048, "auto", 3000), + ("gpt-5.2", 4096, 512, "auto", 615), + ("gpt-4.1-mini-2025-04-14", 1024, 1024, "low", 1659), + ("claude-haiku-5-5", 200, 200, "auto", 64), + ("anthropic/claude-haiku-5-5", 1024, 1024, "auto", 1369), + ("claude-haiku-5-5", 1920, 1080, "low", 2691), + ("claude-haiku-4-5", 1920, 1080, "auto", 1560), + ("claude-haiku-4-5", 1075, 1520, "auto", 1551), + ("claude-haiku-5-5", 3840, 2160, "auto", 4784), + ("claude-haiku-5-5", 8000, 100, "auto", 184), + ], +) +def test_image_estimate_uses_model_resolution_and_detail( + model: str, width: int, height: int, detail: str, expected: int +) -> None: + assert estimate_image_tokens(model, _png_data_url(width, height), detail) == expected + + +@pytest.mark.parametrize( + ("models", "fallback", "expected_model"), + [ + (("azure/deployment-prod",), "gpt-4o", "gpt-4o"), + (("openai/gpt-4o",), "gpt-4o-mini", "gpt-4o"), + (("azure/deployment-prod", "openai/gpt-4o"), "gpt-4o-mini", "gpt-4o-mini"), + (("openai/gpt-4o", "openrouter/openai/gpt-4o-mini"), None, "gpt-4o-mini"), + (("azure_ai/openai/gpt-6-astra",), None, "gpt-6-astra"), + (("azure/gpt-41",), None, "gpt-4.1"), + (("azure/eu/gpt-4o-mini",), None, "gpt-4o-mini"), + (("unknown", "gpt-4o-mini"), None, "gpt-4o-mini"), + ((), "gpt-6-astra", "gpt-6-astra"), + (("other/gpt-4o-mini",), None, "unregistered-model"), + (("azure/opaque/gpt-4o-mini",), None, "unregistered-model"), + (("bedrock/anthropic.claude-haiku-5-5",), None, "claude-haiku-5-5"), + (("bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0",), None, "claude-haiku-4-5"), + (("vertex_ai/claude-haiku-5-5@default",), None, "claude-haiku-5-5"), + (("azure_ai/claude-haiku-5-5",), None, "claude-haiku-5-5"), + (("anthropic/claude-haiku-5-5", "openai/gpt-4o"), "gpt-4o-mini", "claude-haiku-5-5"), + (("other/claude-haiku-5-5",), None, "unregistered-model"), + ], +) +def test_image_estimate_prefers_each_backend_policy_over_the_requested_model( + models: tuple[str, ...], fallback: str | None, expected_model: str +) -> None: + data: Final = _png_data_url(1024, 1024) + assert estimate_image_tokens(models, data, "high", fallback_model=fallback) == estimate_image_tokens( + expected_model, data, "high" + ) + + +@pytest.mark.parametrize( + ("model", "detail"), + [ + ("unknown", "high"), + ("claude-unregistered", "auto"), + ("gpt-4o", "original"), + ("gpt-5.2", "original"), + ("gpt-6-astra", "invalid"), + ], +) +def test_image_estimate_preserves_default_when_policy_is_unavailable(model: str, detail: str) -> None: + assert ( + estimate_image_tokens(model, _png_data_url(1024, 1024), detail) == litellm.constants.DEFAULT_IMAGE_TOKEN_COUNT + ) + + +@pytest.mark.parametrize( + "data", + [ + "https://image.invalid/photo.png", + "data:image/png;base64,invalid!", + "data:application/octet-stream;base64,bm90LWFuLWltYWdl", + "data:image/jpeg;base64,/9j/4AAQ", + "data:image/webp;base64,UklGRgAAAABXRUJQ", + _png_data_url(0, 1024), + ], +) +# Accepted patch limits verified 2026-10-28: https://developers.openai.com/api/docs/guides/images-vision +# and https://platform.claude.com/docs/en/build-with-claude/vision +@pytest.mark.parametrize( + ("model", "detail", "bound"), + [ + ("gpt-4o", "high", high_detail_image_token_upper_bound()), + ("gpt-6-astra", "high", 3000), + ("gpt-6-astra", "auto", 36000), + ("claude-haiku-5-5", "auto", 4784), + ("claude-haiku-4-5", "auto", 1568), + ], +) +def test_image_estimate_uses_bounded_fallback_without_fetching_media( + data: str, model: str, detail: str, bound: int +) -> None: + with patch("httpx.Client.send") as send: + estimate: Final = estimate_image_tokens(model, data, detail) + send.assert_not_called() + assert estimate == bound + assert estimate > estimate_image_tokens(model, _png_data_url(1024, 1024), detail) + + +@pytest.mark.parametrize("container", ["PNG", "JPEG", "GIF", "WEBP"]) +def test_image_estimate_depends_on_dimensions_and_not_compressed_payload_size(container: str) -> None: + output: Final = io.BytesIO() + Image.new("RGB", (1024, 1024)).save(output, format=container) + data: Final = "data:application/octet-stream;base64," + base64.b64encode(output.getvalue()).decode() + assert estimate_image_tokens("gpt-6-astra", data) == estimate_image_tokens("gpt-6-astra", _png_data_url(1024, 1024)) + + +def test_image_estimate_bounds_header_decoding_for_large_inline_payloads() -> None: + header: Final = base64.b64decode(_png_data_url(1024, 1024).split(",", 1)[1]) + data: Final = "data:image/png;base64," + base64.b64encode(header + b"\0" * 4_000_000).decode() + with patch("base64.b64decode", wraps=base64.b64decode) as decode_header: + estimate: Final = estimate_image_tokens("gpt-6-astra", data) + assert estimate == estimate_image_tokens("gpt-6-astra", _png_data_url(1024, 1024)) + assert decode_header.call_count == 1 + assert len(decode_header.call_args.args[0]) <= 65_536 + + HUB_SETUP_SCRIPT: Final = """ import json import sys diff --git a/tests/unit/litellm_core_utils/test_tokenizer.py b/tests/unit/litellm_core_utils/test_tokenizer.py index 61e797d7166..54185c08c10 100644 --- a/tests/unit/litellm_core_utils/test_tokenizer.py +++ b/tests/unit/litellm_core_utils/test_tokenizer.py @@ -241,20 +241,14 @@ print("compatible") @pytest.mark.parametrize("rust", (None, "0", "1")) -def test_tokenization_without_native_extension_stays_offline(tmp_path: Path, rust: str | None) -> None: +def test_tokenization_stays_offline(tmp_path: Path, rust: str | None) -> None: script: Final = """ -import importlib.abc import sys sys.path.insert(0, sys.argv[1]) def reject_network(event, args): if event == "socket.connect": raise AssertionError("tokenizer attempted a network connection") sys.addaudithook(reject_network) -class Block(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == "litellm.rust_bridge._native": - raise ImportError("native extension is unavailable") -sys.meta_path.insert(0, Block()) import litellm from litellm.rust_bridge.tokenizer import get_encoding import tiktoken diff --git a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py index f3b787beda1..d4490927e29 100644 --- a/tests/unit/litellm_proxy_extras/test_migration_data_checks.py +++ b/tests/unit/litellm_proxy_extras/test_migration_data_checks.py @@ -7,7 +7,9 @@ from litellm_proxy_extras.migration_data_checks import ( MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR, BlockedChange, Check, + ColumnType, CouldNotVerify, + _count_query, btree_max_entry_bytes, checks_for_index, checks_for_migration, @@ -113,6 +115,19 @@ def test_the_btree_entry_limit_on_the_default_page_is_the_number_in_postgres_err assert btree_max_entry_bytes(8192) == 2704 +def test_the_entry_size_query_sizes_a_text_column_by_its_stored_bytes_when_postgres_compressed_it() -> None: + """Postgres compresses a long value before it checks the entry size, so a repetitive + 3000-byte key indexes fine. The query has to size it as stored, never by its length alone.""" + (check,) = checks_for_index( + "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime") + ) + types: Final = (ColumnType("text", -1, "i", -1), ColumnType("timestamp", 8, "d", 3)) + query: Final = _count_query(check, "public", types, 2704).as_string(None) + stored: Final = 'least(octet_length("api_key"), pg_column_size("api_key"))' + assert stored in query + assert 'octet_length("api_key")' not in query.replace(stored, ""), query + + def test_the_report_names_the_rows_the_statement_and_the_fix_and_says_nothing_was_applied() -> None: (check,) = checks_for_index( "LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", ("api_key", "startTime") diff --git a/tests/unit/llms/aiohttp_openai/chat/test_transformation.py b/tests/unit/llms/aiohttp_openai/chat/test_transformation.py index 253f9dbe746..ae242e69a9a 100644 --- a/tests/unit/llms/aiohttp_openai/chat/test_transformation.py +++ b/tests/unit/llms/aiohttp_openai/chat/test_transformation.py @@ -1,15 +1,16 @@ +from typing import Final from unittest.mock import AsyncMock, Mock import pytest from aiohttp import ClientResponse -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError from litellm.llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig from litellm.types.utils import ModelResponse -async def _transform(body: object) -> ModelResponse: - raw_response = Mock(spec=ClientResponse) +async def _transform(body: object, *, requested_tier: str | None = None) -> ModelResponse: + raw_response: Final = Mock(spec=ClientResponse) raw_response.json = AsyncMock(return_value=body) return await AiohttpOpenAIChatConfig().transform_response( model="gpt-4o", @@ -18,7 +19,7 @@ async def _transform(body: object) -> ModelResponse: logging_obj=Mock(), request_data={}, messages=[{"role": "user", "content": "Hello"}], - optional_params={}, + optional_params={"service_tier": requested_tier} if requested_tier is not None else {}, litellm_params={}, encoding=None, ) @@ -99,3 +100,20 @@ async def test_transform_response_rejects_choices_that_are_not_a_list_of_objects async def test_transform_response_rejects_a_choice_whose_message_is_not_an_object(): with pytest.raises(ValidationError, match="validation error for Choices"): await _transform({"choices": [{"message": "text"}]}) + + +@pytest.mark.parametrize("tier", ("flex", "priority", "default", "custom-tier", None)) +async def test_transform_response_serializes_served_tier_instead_of_requested_tier(tier: str | None) -> None: + response: Final = await _transform( + {"id": "chatcmpl-1", "choices": [{}], "service_tier": tier}, requested_tier="auto" + ) + serialized: Final = TypeAdapter(dict[str, object]).validate_json(response.model_dump_json()) + + assert serialized["service_tier"] == tier + + +async def test_transform_response_does_not_invent_a_served_tier_when_omitted() -> None: + response: Final = await _transform({"id": "chatcmpl-1", "choices": [{}]}, requested_tier="flex") + serialized: Final = TypeAdapter(dict[str, object]).validate_json(response.model_dump_json()) + + assert "service_tier" not in serialized diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py index 1f4a5aaa256..355d407001c 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_handler.py @@ -15,7 +15,7 @@ import litellm from litellm._uuid import uuid from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt -from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call +from litellm.llms.anthropic.chat.handler import ModelResponseIterator, _envelope_verdict, make_call, make_sync_call from litellm.llms.base_llm.base_utils import response_format_schema, unwrap_response_format_envelope from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.types.llms.openai import ( @@ -27,24 +27,34 @@ from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream from tests.unit.llms.anthropic.anthropic_docs_refusal import ANTHROPIC_DOCS_REFUSAL_MESSAGE +def _text_delta_stream(text: str) -> bytes: + events: Final = ( + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}), + ) + return "".join( + f"event: {name}\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n" for name, payload in events + ).encode() + + +def _streamed_text(chunks: list[object]) -> str: + return "".join(chunk.choices[0].delta.content or "" for chunk in chunks if isinstance(chunk, ModelResponseStream)) + + @pytest.mark.asyncio async def test_make_call_passes_logging_obj_to_client_post(): """make_call must pass logging_obj to client.post so track_llm_api_timing can set llm_api_duration_ms for litellm_overhead_time_ms.""" mock_client = AsyncMock() - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock( - return_value=iter([b'data: {"type":"message_start"}\n', b'data: {"type":"message_delta"}\n']) - ) - mock_client.post.return_value = mock_response + mock_client.post.return_value = httpx.Response(200, stream=httpx.ByteStream(_text_delta_stream("Hi there"))) logging_obj = MagicMock() - await make_call( + completion_stream, _ = await make_call( client=mock_client, api_base="https://api.anthropic.com/v1/messages", headers={}, data="{}", - model="claude-3-5-haiku", + model="claude-haiku-5-5", messages=[{"role": "user", "content": "Hi"}], logging_obj=logging_obj, timeout=60.0, @@ -54,6 +64,39 @@ async def test_make_call_passes_logging_obj_to_client_post(): mock_client.post.assert_called_once() call_kwargs = mock_client.post.call_args[1] assert call_kwargs.get("logging_obj") is logging_obj + assert _streamed_text([chunk async for chunk in completion_stream]) == "Hi there" + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +async def _streamed_chunks(sync_stream: bool, text: str) -> list[object]: + transport: Final = httpx.MockTransport( + lambda request: httpx.Response(200, stream=httpx.ByteStream(_text_delta_stream(text))) + ) + call_args: Final = { + "api_base": "https://api.anthropic.com/v1/messages", + "headers": {}, + "data": "{}", + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "Hi"}], + "logging_obj": MagicMock(), + "timeout": 60.0, + "json_mode": False, + } + if sync_stream: + sync_iterator, _ = make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args) + return list(sync_iterator) + async_iterator, _ = await make_call(client=AsyncHTTPHandler(transport=transport), **call_args) + return [chunk async for chunk in async_iterator] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_streamed_text_keeps_unicode_line_separators(sync_stream: bool) -> None: + chunks: Final = await _streamed_chunks(sync_stream, _LINE_SEPARATED_TEXT) + + assert _streamed_text(chunks) == _LINE_SEPARATED_TEXT def test_anthropic_completion_does_not_send_deployment_default_limits(): diff --git a/tests/unit/llms/anthropic/messages/context_management/test_compact.py b/tests/unit/llms/anthropic/messages/context_management/test_compact.py index 159c3586c53..dff931a6074 100644 --- a/tests/unit/llms/anthropic/messages/context_management/test_compact.py +++ b/tests/unit/llms/anthropic/messages/context_management/test_compact.py @@ -1523,6 +1523,7 @@ async def test_summary_model_denied_when_user_over_model_budget(): user_id="user-over-budget", user_model_max_budget={"claude-haiku-4-5": {"budget_limit": 5}}, model="claude-haiku-4-5", + team_id=None, ) import inspect diff --git a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py index 5af8e69a0ea..4626f76c8fc 100644 --- a/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py +++ b/tests/unit/llms/anthropic/messages/responses_adapters/test_responses_adapters_streaming_iterator.py @@ -4,20 +4,25 @@ Tests for AnthropicResponsesStreamWrapper """ import asyncio +import datetime import json import os import sys +from collections.abc import AsyncIterator, Awaitable, Callable from types import SimpleNamespace +import anyio import pytest sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) import litellm from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.prompt_templates.common_utils import ( encrypted_reasoning_signature, ) +from litellm.llms.anthropic.messages.native.streaming_iterator import aclose_superseded_attempt from litellm.llms.anthropic.messages.native.utils import INCOMPLETE_STREAM_ERROR_MESSAGE from litellm.llms.anthropic.messages.responses_adapters.streaming_iterator import ( AnthropicResponsesStreamWrapper, @@ -622,3 +627,194 @@ class TestUpstreamFailureEndsStreamWithErrorEvent: chunks = _collect(iter(())) assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] assert chunks[1]["error"] == {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE} + + +class _FailureHookRecorder: + def __init__(self) -> None: + self.received: tuple[Exception, ...] = () + + async def __call__(self, failure: Exception) -> None: + self.received = (*self.received, failure) + + +def _logging_obj_with_failure_hook(on_failure: Callable[[Exception], Awaitable[None]]) -> LitellmLogging: + logging_obj = LitellmLogging( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="lit-6871-test", + function_id="lit-6871-test", + ) + logging_obj._on_detached_stream_failure = on_failure + return logging_obj + + +def _collect_with_logging( + stream: AsyncIterator[dict[str, object]], logging_obj: LitellmLogging +) -> list[dict[str, object]]: + async def _run() -> list[dict[str, object]]: + wrapper = AnthropicResponsesStreamWrapper(responses_stream=stream, model="m", litellm_logging_obj=logging_obj) + return [chunk async for chunk in wrapper] + + return asyncio.run(_run()) + + +class _DeliveryCountingConsumer: + def __init__(self) -> None: + self.delivered: tuple[dict[str, object], ...] = () + self.delivered_when_hook_fired: tuple[int, ...] = () + + async def hook(self, failure: Exception) -> None: + self.delivered_when_hook_fired = (*self.delivered_when_hook_fired, len(self.delivered)) + + async def drain(self, wrapper: AnthropicResponsesStreamWrapper) -> None: + async for chunk in wrapper: + self.delivered = (*self.delivered, chunk) + + +def _failing_after_one_delta() -> AsyncIterator[dict[str, object]]: + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} + yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} + raise litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m") + + return _gen() + + +class TestUpstreamFailureReachesTheProxyFailureHook: + def test_raised_upstream_failure_reaches_the_hook_unwrapped_once_the_error_event_was_consumed(self): + rate_limit = litellm.RateLimitError(message="You have no credits remaining.", llm_provider="openai", model="m") + wrapped = MidStreamFallbackError( + message=str(rate_limit), + model="m", + llm_provider="openai", + original_exception=rate_limit, + is_pre_first_chunk=False, + ) + recorder = _FailureHookRecorder() + + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + yield {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_1"}} + yield {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "Hi"} + raise wrapped + + chunks = _collect_with_logging(_gen(), _logging_obj_with_failure_hook(recorder)) + assert [chunk["type"] for chunk in chunks] == [ + "message_start", + "content_block_start", + "content_block_delta", + "error", + ] + assert recorder.received == (rate_limit,) + + def test_response_failed_event_reaches_the_hook_as_an_exception_carrying_the_upstream_status(self): + recorder = _FailureHookRecorder() + failed = SimpleNamespace( + status="failed", + output=[], + usage=None, + error={"code": "rate_limit_exceeded", "message": "Rate limit reached for gpt-5.5, try again in 20s."}, + ) + + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + yield {"type": "response.failed", "response": failed} + + chunks = _collect_with_logging(_gen(), _logging_obj_with_failure_hook(recorder)) + assert chunks[-1]["type"] == "error" + (failure,) = recorder.received + assert isinstance(failure, litellm.RateLimitError) + assert failure.status_code == 429 + assert "Rate limit reached for gpt-5.5, try again in 20s." in failure.message + + def test_a_raising_hook_never_hides_the_error_event_from_the_client(self): + async def _exploding_hook(failure: Exception) -> None: + raise RuntimeError("spend writer is down") + + logging_obj = _logging_obj_with_failure_hook(_exploding_hook) + + async def _gen() -> AsyncIterator[dict[str, object]]: + yield {"type": "response.created"} + raise ConnectionResetError("Response payload is not completed") + + chunks = _collect_with_logging(_gen(), logging_obj) + assert [chunk["type"] for chunk in chunks] == ["message_start", "error"] + assert chunks[-1]["error"] == {"type": "api_error", "message": "Response payload is not completed"} + + def test_the_hook_fires_only_after_the_consumer_took_the_error_event(self): + consumer = _DeliveryCountingConsumer() + + async def _run() -> None: + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=_failing_after_one_delta(), + model="m", + litellm_logging_obj=_logging_obj_with_failure_hook(consumer.hook), + ) + await consumer.drain(wrapper) + with pytest.raises(StopAsyncIteration): + await wrapper.__anext__() + + asyncio.run(_run()) + assert [chunk["type"] for chunk in consumer.delivered] == [ + "message_start", + "content_block_start", + "content_block_delta", + "error", + ] + assert consumer.delivered_when_hook_fired == (len(consumer.delivered),) + + @pytest.mark.parametrize( + ("close", "expected_reports"), + [ + pytest.param(lambda sse: sse.aclose(), 1, id="client_closes"), + pytest.param(aclose_superseded_attempt, 0, id="fallback_takes_over"), + ], + ) + def test_closing_the_stream_on_the_error_event_reports_unless_a_fallback_takes_over(self, close, expected_reports): + recorder = _FailureHookRecorder() + + async def _close_on_the_error_event() -> None: + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(recorder) + ) + sse = wrapper.async_anthropic_sse_wrapper() + async for frame in sse: + if frame.startswith(b"event: error\n"): + await close(sse) + return + pytest.fail("the stream ended without an error event") + + asyncio.run(_close_on_the_error_event()) + assert len(recorder.received) == expected_reports + + def test_a_client_disconnect_during_the_failure_hook_lets_the_hook_finish(self): + hook_started = anyio.Event() + hook_released = anyio.Event() + finished: list[Exception] = [] + + async def _held_hook(failure: Exception) -> None: + hook_started.set() + await hook_released.wait() + finished.append(failure) + + async def _disconnect_mid_hook() -> None: + wrapper = AnthropicResponsesStreamWrapper( + responses_stream=_failing_after_one_delta(), model="m", litellm_logging_obj=_logging_obj_with_failure_hook(_held_hook) + ) + + async def _consume() -> None: + async for _ in wrapper.async_anthropic_sse_wrapper(): + pass + + async with anyio.create_task_group() as task_group: + task_group.start_soon(_consume) + await hook_started.wait() + task_group.cancel_scope.cancel() + hook_released.set() + + asyncio.run(_disconnect_mid_hook()) + assert len(finished) == 1 diff --git a/tests/unit/llms/anthropic/messages/test_stream_assembly.py b/tests/unit/llms/anthropic/messages/test_stream_assembly.py index f7d6b8d4f8c..4993d2fa9bf 100644 --- a/tests/unit/llms/anthropic/messages/test_stream_assembly.py +++ b/tests/unit/llms/anthropic/messages/test_stream_assembly.py @@ -144,6 +144,30 @@ def test_anthropic_message_delta_fields_reads_the_fields_the_openai_shape_drops( assert dict(anthropic_message_delta_fields(GOOGLE_STREAM)) == {} +def test_anthropic_message_delta_fields_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + safeguard_results: Final = [ + { + "type": "dangerous_tool_use", + "status": { + "type": "available", + "tool_uses": { + "toolu_01": {"type": "evaluated", "outcome": "flagged", "explanation": "deletes\u2028the\u2029root"} + }, + }, + } + ] + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 2}, + } + raw_message_delta_frame: Final = f"event: message_delta\ndata: {json.dumps(message_delta, ensure_ascii=False)}\n\n" + + fields: Final = anthropic_message_delta_fields((*ANTHROPIC_STREAM[:5], raw_message_delta_frame.encode())) + + assert dict(fields) == {"safeguard_results": safeguard_results} + + @pytest.mark.parametrize( "chunks", [GOOGLE_STREAM, INVALID_UTF8, (ANTHROPIC_ERROR,), ()], diff --git a/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py index eabd5c8427d..83e01fce642 100644 --- a/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/unit/llms/azure/image_generation/test_azure_image_generation_init.py @@ -20,27 +20,9 @@ from litellm.llms.azure.image_generation.http_utils import ( azure_deployment_image_generation_json_body, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.llms.azure.image_generation import ( - AzureDallE3ImageGenerationConfig, - get_azure_image_generation_config, -) from litellm.utils import get_optional_params_image_gen -@pytest.mark.parametrize( - "received_model, expected_config", - [ - ("dall-e-3", AzureDallE3ImageGenerationConfig), - ("dalle-3", AzureDallE3ImageGenerationConfig), - ("openai_dall_e_3", AzureDallE3ImageGenerationConfig), - ], -) -def test_azure_image_generation_config(received_model, expected_config): - assert isinstance( - get_azure_image_generation_config(received_model), expected_config - ) - - def test_azure_deployment_image_generation_json_body(): """Deployment-scoped Azure image URL must not send ``model`` in JSON.""" api = ( diff --git a/tests/unit/llms/azure_ai/agents/test_handler.py b/tests/unit/llms/azure_ai/agents/test_handler.py index 71466bb172c..14b6094747f 100644 --- a/tests/unit/llms/azure_ai/agents/test_handler.py +++ b/tests/unit/llms/azure_ai/agents/test_handler.py @@ -173,7 +173,7 @@ async def test_handler_acompletion_with_thread_id_reuses_that_thread() -> None: def _sse(event: str, data: dict[str, object] | str) -> str: - return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data)}\n\n" + return f"event: {event}\ndata: {data if isinstance(data, str) else json.dumps(data, ensure_ascii=False)}\n\n" def _message_delta(text: str) -> dict[str, object]: @@ -222,3 +222,24 @@ async def test_acompletion_stream_yields_agent_message_deltas_then_stop(respx_mo "thread": {"messages": [{"role": "user", "content": "Hi Agent, what is 10 + 5?"}]}, } assert chunks == [("10 + 5", None), (" = 15", None), (None, "stop")] + + +@pytest.mark.asyncio +async def test_acompletion_stream_keeps_unicode_line_separators_inside_a_delta(respx_mock: respx.MockRouter) -> None: + separated_text: Final = "before\u2028between\u2029after" + _ = respx_mock.post(f"{API_BASE}/threads/runs").respond( + 200, + text=_sse("thread.message.delta", _message_delta(separated_text)) + _sse("done", "[DONE]"), + headers={"content-type": "text/event-stream"}, + ) + + response: Final = await litellm.acompletion( + model=f"azure_ai/agents/{AGENT_ID}", + messages=[{"role": "user", "content": "Say it with separators"}], + api_base=API_BASE, + api_key="azure-ad-token", + stream=True, + ) + chunks: Final = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response] + + assert chunks == [(separated_text, None), (None, "stop")] diff --git a/tests/unit/llms/azure_ai/agents/test_transformation.py b/tests/unit/llms/azure_ai/agents/test_transformation.py index 10676438858..413908c8313 100644 --- a/tests/unit/llms/azure_ai/agents/test_transformation.py +++ b/tests/unit/llms/azure_ai/agents/test_transformation.py @@ -1,7 +1,7 @@ import json from typing import Final -from unittest.mock import MagicMock +import httpx import pytest from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig @@ -428,15 +428,10 @@ async def test_azure_ai_agents_streaming_accumulates_annotations_from_multiple_t "data: [DONE]", ] - async def mock_aiter_lines(): - for line in sse_lines: - yield line - - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines()) + sse_response = httpx.Response(200, stream=httpx.ByteStream("\n".join(sse_lines).encode())) chunks = [] - async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"): + async for chunk in handler._process_sse_stream(sse_response, "azure_ai/agents/asst_123"): chunks.append(chunk) final_chunk = chunks[-1] @@ -498,15 +493,10 @@ async def test_azure_ai_agents_streaming_annotations_from_completed_message(): "data: [DONE]", ] - async def mock_aiter_lines(): - for line in sse_lines: - yield line - - mock_response = MagicMock() - mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines()) + sse_response = httpx.Response(200, stream=httpx.ByteStream("\n".join(sse_lines).encode())) chunks = [] - async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"): + async for chunk in handler._process_sse_stream(sse_response, "azure_ai/agents/asst_123"): chunks.append(chunk) assert len(chunks) >= 1 diff --git a/tests/unit/llms/base_llm/guardrail_translation/test_utils.py b/tests/unit/llms/base_llm/guardrail_translation/test_utils.py index 65208a2bd3b..acd50616f90 100644 --- a/tests/unit/llms/base_llm/guardrail_translation/test_utils.py +++ b/tests/unit/llms/base_llm/guardrail_translation/test_utils.py @@ -1,9 +1,11 @@ +import json from typing import Final import pytest import litellm from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_response_usage, effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, ) @@ -32,3 +34,18 @@ def test_a_per_guardrail_skip_flag_wins_over_the_global_setting( assert effective_skip_system_message_for_guardrail(guardrail) is expected assert effective_skip_tool_message_for_guardrail(guardrail) is expected + + +def test_blocked_response_usage_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + message_start: Final = {"type": "message_start", "message": {"usage": {"input_tokens": 11, "output_tokens": 1}}} + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "stop_sequence", "stop_sequence": "END\u2028OF\u2029TURN"}, + "usage": {"output_tokens": 7}, + } + stream: Final = "".join( + f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n" + for event in (message_start, message_delta) + ).encode() + + assert blocked_response_usage([stream]) == {"input_tokens": 11, "output_tokens": 7} diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index 1890c68b30d..c36644ba536 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -242,6 +242,33 @@ def test_transform_response_hands_json_mode_to_nova(): assert json.loads(result.choices[0].message.content) == {"city": "Paris", "temperature": 21} +@pytest.mark.parametrize(("finish_reason", "expected"), [("MAX_TOKENS", "length"), (None, "stop")]) +def test_transform_response_accepts_null_cohere_finish_reason(finish_reason: str | None, expected: str) -> None: + from litellm.types.utils import ModelResponse + + raw_response: Final = httpx.Response( + 200, + json={"generations": [{"text": "Hello there", "finish_reason": finish_reason}]}, + request=httpx.Request("POST", "https://bedrock"), + ) + + result: Final = AmazonInvokeConfig().transform_response( + model="invoke/cohere.command-text-v14", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + api_key=None, + ) + + assert result.choices[0].message.content == "Hello there" + assert result.choices[0].finish_reason == expected + + def _stream_invoke_completion_with_spied_client(**kwargs: object) -> tuple[MagicMock, MagicMock]: mock_response = MagicMock() mock_response.status_code = 200 diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py index d90e9933b88..5b3cb2e0c92 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_twelvelabs_pegasus_transformation.py @@ -1,6 +1,13 @@ +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest + from litellm.llms.bedrock.chat.invoke_transformations.amazon_twelvelabs_pegasus_transformation import ( AmazonTwelveLabsPegasusConfig, ) +from litellm.types.utils import ModelResponse def _make_messages() -> list[dict]: @@ -84,3 +91,24 @@ def test_transform_request_includes_s3_media(): s3_location = request["mediaSource"]["s3Location"] assert s3_location["uri"] == "s3://test-bucket/video.mp4" assert s3_location["bucketOwner"] == "123456789012" + + +@pytest.mark.parametrize( + ("finish_reason", "expected"), + [("length", "length"), (None, "stop")], +) +def test_transform_response_accepts_null_finish_reason(finish_reason: str | None, expected: str) -> None: + result: Final = AmazonTwelveLabsPegasusConfig().transform_response( + model="twelvelabs.pegasus-1-2-v1:0", + raw_response=httpx.Response(200, json={"message": "A cat jumps.", "finishReason": finish_reason}), + model_response=ModelResponse(), + logging_obj=Mock(), + request_data={}, + messages=_make_messages(), + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result.choices[0].message.content == "A cat jumps." + assert result.choices[0].finish_reason == expected diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 7e05ef500f5..9d9e39df944 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -4458,6 +4458,20 @@ class TestMalformedBatchJsonl: _VALID_BATCH_RECORD: Final = json.loads(_VALID_BATCH_LINE) +def test_file_content_conversion_keeps_a_line_whose_json_holds_unicode_line_separators() -> None: + text: Final = "Hello\u2028from\u2029batch" + body: Final = {**_VALID_BATCH_RECORD["body"], "messages": [{"role": "user", "content": text}]} + record: Final = {**_VALID_BATCH_RECORD, "body": body} + file_content: Final = f"{json.dumps(record, ensure_ascii=False)}\n".encode() + + bedrock_jsonl, _ = BedrockJsonlFilesTransformation().transform_openai_file_content_to_bedrock_file_content( + file_content + ) + + (bedrock_record,) = (json.loads(line) for line in bedrock_jsonl.split("\n")) + assert bedrock_record["modelInput"]["messages"] == [{"role": "user", "content": [{"text": text}]}] + + def _batch_line(body: object) -> str: return json.dumps({**_VALID_BATCH_RECORD, "body": body}) diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 5aaee1f0ee4..aec1672b335 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -18,7 +18,7 @@ import types from collections.abc import Callable, Mapping, Sequence from datetime import datetime from itertools import chain -from typing import ClassVar, Literal, TypeAlias +from typing import ClassVar, Final, Literal, TypeAlias from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -60,6 +60,7 @@ from litellm.llms.bedrock.passthrough.guardrail_translation.handler import ( BedrockPassthroughGuardrailHandler, _extract_converse_texts, _is_converse_endpoint, + _sse_events, _write_back_texts, ) @@ -3840,3 +3841,14 @@ class TestInvokeStreamGuardrailRelay: assert received_key is not None and received_key.texts == ("In Paris",) assert rewritten_key == received_key + + +def test_sse_events_keeps_a_data_line_whose_json_holds_unicode_line_separators() -> None: + event: Final = { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Card\u2028on\u2029file"}, + } + sse: Final = f"event: content_block_delta\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + + assert _sse_events(sse) == (event,) diff --git a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py index 664dbf87fe9..01cca1fc270 100644 --- a/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py +++ b/tests/unit/llms/bedrock/search/test_agentcore_search_transformation.py @@ -8,6 +8,7 @@ the sharded CI (coverage collection runs against this tree). import json import os +from typing import Final import httpx import pytest @@ -682,3 +683,17 @@ def test_transform_search_response_ignores_text_blocks_without_result_objects(bl body = {"result": {"content": [{"type": "text", "text": block_text}]}} assert _transform_text(json.dumps(body)).results == [] + + +def test_transform_search_response_keeps_a_result_whose_sse_json_holds_unicode_line_separators() -> None: + snippet: Final = "Snippet\u2028for\u2029result" + results: Final = [{"title": "Result", "url": "https://example.com/1", "text": snippet}] + body: Final = { + "jsonrpc": "2.0", + "id": 1, + "result": {"content": [{"type": "text", "text": json.dumps(results, ensure_ascii=False)}]}, + } + + response: Final = _transform_text(f"event: message\ndata: {json.dumps(body, ensure_ascii=False)}\n\n") + + assert _as_tuples(response) == [("Result", "https://example.com/1", snippet, None, None)] diff --git a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py index a6f7d613f28..635363e77ea 100644 --- a/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py +++ b/tests/unit/llms/chatgpt/responses/test_chatgpt_responses_transformation.py @@ -587,6 +587,37 @@ class TestChatGPTResponsesMultiAccount: assert headers_default["ChatGPT-Account-Id"] == "acct-default" +def test_chatgpt_non_stream_sse_keeps_an_output_item_whose_text_holds_unicode_line_separators() -> None: + text: Final = "Hello\u2028from\u2029stream" + output_item_done: Final = { + "type": "response.output_item.done", + "output_index": 0, + "item": {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}, + } + completed: Final = { + "type": "response.completed", + "response": { + "id": "resp_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-6.1-sol", + "output": [], + }, + } + sse_body: Final = "".join( + f"data: {json.dumps(event, ensure_ascii=False)}\n" for event in (output_item_done, completed) + ) + + parsed: Final = ChatGPTResponsesAPIConfig().transform_response_api_response( + model="chatgpt/gpt-6.1-sol", + raw_response=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=sse_body), + logging_obj=MagicMock(), + ) + + assert parsed.output_text == text + + @pytest.mark.parametrize("call", ["responses", "completion"]) def test_prompt_cache_key_keeps_one_session_header_across_requests( call: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/llms/cloudflare/test_cloudflare_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py index 3a45c83556c..ccbfc02f0e6 100644 --- a/tests/unit/llms/cloudflare/test_cloudflare_transformation.py +++ b/tests/unit/llms/cloudflare/test_cloudflare_transformation.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import Callable, Iterator +from collections.abc import AsyncIterator, Callable, Iterator from typing import Final from unittest.mock import MagicMock @@ -380,6 +380,35 @@ def test_completion_cloudflare_stream(sync_mode: bool, _cloudflare_httpx_transpo assert "language" in content.lower() +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +async def test_cloudflare_stream_keeps_unicode_line_separators_inside_content( + sync_mode: bool, _cloudflare_httpx_transport: None +) -> None: + streamed_text: Final = ("alpha\u2028", "beta\u2029", "gamma") + base: Final = {"id": "chatcmpl-cf", "object": "chat.completion.chunk", "created": 1, "model": "@cf/glm"} + body: Final = "".join( + f"data: {json.dumps({**base, 'choices': [{'index': 0, 'delta': {'content': text}}]}, ensure_ascii=False)}\n\n" + for text in streamed_text + ) + with respx.mock(assert_all_called=True) as api: + api.post(f"{FAKE_API_BASE}/chat/completions").respond( + 200, headers={"content-type": "text/event-stream"}, content=(body + "data: [DONE]\n\n").encode() + ) + request: Final = { + "model": "cloudflare/@cf/glm", + "messages": [{"role": "user", "content": "repeat"}], + "stream": True, + "api_base": FAKE_API_BASE, + "api_key": FAKE_API_KEY, + } + chunks: Final = ( + tuple(completion(**request)) if sync_mode else tuple([c async for c in await acompletion(**request)]) + ) + + assert "".join(c.choices[0].delta.content or "" for c in chunks) == "".join(streamed_text) + + def _make_mock_response(json_data: Dict[str, Any]) -> MagicMock: mock = MagicMock(spec=httpx.Response) mock.status_code = 200 @@ -397,13 +426,13 @@ def test_acompletion_cloudflare_stream(sync_mode, monkeypatch): if sync_mode: - def _iter_lines(): + def _iter_text() -> Iterator[str]: for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" + yield f"data: {chunk}\n\n" + yield "data: [DONE]\n\n" mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() + mock_resp.iter_text.return_value = _iter_text() mock_resp.status_code = 200 mock_resp.headers = {"content-type": "text/event-stream"} @@ -420,13 +449,13 @@ def test_acompletion_cloudflare_stream(sync_mode, monkeypatch): mock_post.assert_called_once() else: - async def _aiter_lines(): + async def _aiter_text() -> AsyncIterator[str]: for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" + yield f"data: {chunk}\n\n" + yield "data: [DONE]\n\n" mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() + mock_resp.aiter_text.return_value = _aiter_text() mock_resp.status_code = 200 mock_resp.headers = {"content-type": "text/event-stream"} diff --git a/tests/unit/llms/codestral/__init__.py b/tests/unit/llms/codestral/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codestral/completion/__init__.py b/tests/unit/llms/codestral/completion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py new file mode 100644 index 00000000000..9ce50d8ceca --- /dev/null +++ b/tests/unit/llms/codestral/completion/test_codestral_completion_handler.py @@ -0,0 +1,50 @@ +import json +from typing import Final + +import httpx +import pytest +import respx + +import litellm + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _fim_stream_body(text: str) -> bytes: + chunks: Final = ( + {"id": "c1", "object": "chat.completion.chunk", "created": 1, "model": "codestral-latest", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": None}]}, + {"id": "c1", "object": "chat.completion.chunk", "created": 1, "model": "codestral-latest", + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}]}, + ) + events: Final = "".join(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" for chunk in chunks) + return f"{events}data: [DONE]\n\n".encode() + + +async def _streamed_texts(sync_stream: bool) -> list[str | None]: + request: Final = { + "model": "text-completion-codestral/codestral-latest", + "prompt": "def f(", + "stream": True, + "api_key": "codestral-key", + "max_retries": 0, + } + if sync_stream: + return [chunk.choices[0].text for chunk in litellm.text_completion(**request)] + return [chunk.choices[0].text async for chunk in await litellm.atext_completion(**request)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_fim_stream_keeps_unicode_line_separators_inside_a_delta( + sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post("https://codestral.mistral.ai/v1/fim/completions").mock( + return_value=httpx.Response(200, content=_fim_stream_body(_LINE_SEPARATED_TEXT)) + ) + + texts: Final = await _streamed_texts(sync_stream) + + assert texts == [_LINE_SEPARATED_TEXT, None] + assert json.loads(route.calls.last.request.content)["stream"] is True diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py index 83c38b65beb..bb66a70da84 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_handler.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_handler.py @@ -1,11 +1,17 @@ +import json +from typing import Final from unittest.mock import AsyncMock, Mock, patch import aiohttp +import httpx import pytest import litellm +from litellm.llms.aiohttp_openai.chat.transformation import AiohttpOpenAIChatConfig from litellm.llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.utils import ModelResponseStream class TestBaseLLMAIOHTTPHandler: @@ -445,3 +451,34 @@ class TestBaseLLMAIOHTTPHandler: # Should use transport, not connector mock_transport.get_valid_client_session.assert_called_once() assert result is mock_session_from_transport + + +def test_make_sync_call_stream_keeps_unicode_line_separators_inside_a_delta() -> None: + text: Final = "before\u2028between\u2029after" + chunk: Final = { + "id": "c1", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.5", + "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}], + } + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n".encode() + transport: Final = httpx.MockTransport(lambda request: httpx.Response(200, stream=httpx.ByteStream(body))) + + completion_stream, _ = BaseLLMAIOHTTPHandler().make_sync_call( + provider_config=AiohttpOpenAIChatConfig(), + api_base="https://example.com/v1/chat/completions", + headers={}, + data={"model": "gpt-5.5", "messages": []}, + model="gpt-5.5", + messages=[], + logging_obj=Mock(), + litellm_params={}, + timeout=5.0, + client=HTTPHandler(client=httpx.Client(transport=transport)), + ) + + contents: Final = [ + chunk.choices[0].delta.content for chunk in completion_stream if isinstance(chunk, ModelResponseStream) + ] + assert contents == [text] diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index c7d5518ef5b..aed808223bd 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -5,6 +5,7 @@ import socket import sys import struct import threading +import weakref from typing import Final import aiohttp @@ -1383,3 +1384,57 @@ def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_ loop.call_soon_threadsafe(loop.stop) thread.join(timeout=10) loop.close() + + +@pytest.mark.asyncio +async def test_a_stream_keeps_reading_after_the_client_that_opened_it_is_collected(): + from aiohttp import web + + release: Final = asyncio.Event() + body: Final = tuple(f"chunk{i}\n".encode() for i in range(4)) + + async def handler(request: web.Request) -> web.StreamResponse: + response: Final = web.StreamResponse() + await response.prepare(request) + await response.write(body[0]) + await release.wait() + for chunk in body[1:]: + await response.write(chunk) + await response.write_eof() + return response + + app: Final = web.Application() + app.router.add_get("/stream", handler) + runner: Final = web.AppRunner(app) + await runner.setup() + sock: Final = socket.socket() + sock.bind(("127.0.0.1", 0)) + port: Final[int] = sock.getsockname()[1] + await web.SockSite(runner, sock).start() + + transport: Final = LiteLLMAiohttpTransport(client=aiohttp.ClientSession) + client: Final = httpx.AsyncClient(transport=transport) + transport_ref: Final = weakref.ref(transport) + try: + response: Final = await client.send(client.build_request("GET", f"http://127.0.0.1:{port}/stream"), stream=True) + session: Final = transport.client + assert isinstance(session, aiohttp.ClientSession) + chunks: Final = response.aiter_bytes() + first: Final = await chunks.__anext__() + + del client, transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + release.set() + + rest: Final = [chunk async for chunk in chunks] + assert first + b"".join(rest) == b"".join(body), (first, rest) + assert not session.closed, "the session was closed under a stream that was still being read" + + del response, chunks + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + assert transport_ref() is None, "the finished stream kept the transport alive" + assert session.closed, "the owned session stayed open after its last stream finished" + finally: + await runner.cleanup() diff --git a/tests/unit/llms/custom_httpx/test_sse_lines.py b/tests/unit/llms/custom_httpx/test_sse_lines.py new file mode 100644 index 00000000000..9bb03f6da98 --- /dev/null +++ b/tests/unit/llms/custom_httpx/test_sse_lines.py @@ -0,0 +1,64 @@ +from collections.abc import AsyncIterator, Sequence +from typing import Final + +import pytest + +from litellm.llms.custom_httpx.sse_lines import ( + aiter_sse_lines, + iter_sse_lines, + split_sse_lines, + split_sse_lines_keepends, +) + +CASES: Final = ( + pytest.param(("data: a\u2028b\r\n\r\n",), ("data: a\u2028b", ""), id="u2028-stays-inside-the-line"), + pytest.param( + ("data: \u0085\u2029\x0b\x0c\x1c\x1d\x1e\n",), + ("data: \u0085\u2029\x0b\x0c\x1c\x1d\x1e",), + id="no-splitlines-only-separator-breaks-a-line", + ), + pytest.param(("a\r\nb\nc\rd\n",), ("a", "b", "c", "d"), id="crlf-lf-and-cr-each-end-a-line"), + pytest.param(("a\r", "\nb\n"), ("a", "b"), id="crlf-split-across-chunks-is-one-break"), + pytest.param(("a\r", "\n\r\n"), ("a", ""), id="crlf-split-then-blank-line"), + pytest.param(("a\r", "", "\nb\n"), ("a", "b"), id="empty-chunk-between-cr-and-lf"), + pytest.param(("a\r", "\r"), ("a", ""), id="cr-then-cr-is-two-breaks"), + pytest.param(("a\r\r\nb\n",), ("a", "", "b"), id="cr-then-crlf-in-one-chunk-is-two-breaks"), + pytest.param(("a\r\nb\r\n\r\n",), ("a", "b", ""), id="crlf-only-chunk"), + pytest.param(("da", "ta: x", "y\n"), ("data: xy",), id="line-spanning-chunks-is-joined"), + pytest.param(("data: x",), ("data: x",), id="unterminated-last-line-is-kept"), + pytest.param(("",), (), id="empty-body"), +) + + +async def _async_chunks(chunks: Sequence[str]) -> AsyncIterator[str]: + for chunk in chunks: + yield chunk + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +def test_iter_sse_lines_splits_only_on_cr_lf_and_crlf(chunks: Sequence[str], expected: Sequence[str]) -> None: + assert tuple(iter_sse_lines(chunks)) == tuple(expected) + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +async def test_aiter_sse_lines_splits_only_on_cr_lf_and_crlf(chunks: Sequence[str], expected: Sequence[str]) -> None: + assert tuple([line async for line in aiter_sse_lines(_async_chunks(chunks))]) == tuple(expected) + + +@pytest.mark.parametrize(("chunks", "expected"), CASES) +def test_split_sse_lines_splits_a_whole_body_only_on_cr_lf_and_crlf( + chunks: Sequence[str], expected: Sequence[str] +) -> None: + assert split_sse_lines("".join(chunks)) == tuple(expected) + + +@pytest.mark.parametrize( + ("text", "expected"), + ( + pytest.param("a\r\nb\rc\nd", ("a\r\n", "b\r", "c\n", "d"), id="each-break-stays-on-its-line"), + pytest.param("data: a\u2028b\u2029c\n\n", ("data: a\u2028b\u2029c\n", "\n"), id="u2028-and-u2029-stay-inside"), + pytest.param("", (), id="empty-body"), + ), +) +def test_split_sse_lines_keepends_splits_only_on_cr_lf_and_crlf(text: str, expected: Sequence[str]) -> None: + assert split_sse_lines_keepends(text) == tuple(expected) diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index b3a224209e9..eb2e1d8e347 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -1,6 +1,6 @@ import copy import json -from collections.abc import Iterator +from collections.abc import AsyncIterator, Iterator from typing import Final from unittest.mock import MagicMock, patch @@ -2160,13 +2160,13 @@ def mock_chat_response_nonanthropic_prompt_caching() -> Dict[str, Any]: def mock_http_handler_chat_streaming_response() -> MagicMock: mock_stream_chunks = mock_chat_streaming_response_chunks() - def mock_iter_lines(): + def mock_iter_text() -> Iterator[str]: for chunk in mock_stream_chunks: for line in chunk.splitlines(): - yield line + yield f"{line}\n" mock_response = MagicMock() - mock_response.iter_lines.side_effect = mock_iter_lines + mock_response.iter_text.side_effect = mock_iter_text mock_response.status_code = 200 return mock_response @@ -2175,13 +2175,13 @@ def mock_http_handler_chat_streaming_response() -> MagicMock: def mock_http_handler_chat_async_streaming_response() -> MagicMock: mock_stream_chunks = mock_chat_streaming_response_chunks() - async def mock_iter_lines(): + async def mock_aiter_text() -> AsyncIterator[str]: for chunk in mock_stream_chunks: for line in chunk.splitlines(): - yield line + yield f"{line}\n" mock_response = MagicMock() - mock_response.aiter_lines.return_value = mock_iter_lines() + mock_response.aiter_text.return_value = mock_aiter_text() mock_response.status_code = 200 return mock_response diff --git a/tests/unit/llms/databricks/test_databricks_cost_calculator.py b/tests/unit/llms/databricks/test_databricks_cost_calculator.py index 7120a130462..c35aa6e4cf0 100644 --- a/tests/unit/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/unit/llms/databricks/test_databricks_cost_calculator.py @@ -33,6 +33,13 @@ PRICE_FIELDS: Final = ( "cache_read_input_token_cost", ) PUBLISHED_DBU_PER_MILLION: Final = { + "databricks/databricks-claude-haiku-5-5": ("1.429", "7.143", "1.786", "0.143"), + "databricks/databricks-claude-sonnet-5-5": ("28.571", "142.857", "35.714", "1.429"), + "databricks/databricks-gpt-6-1-sol": ("28.571", "142.857", "35.714", "1.429"), + "databricks/databricks-gpt-6-1-terra": ("14.286", "71.429", "17.857", "0.714"), + "databricks/databricks-gpt-6-astra": ("142.857", "714.286", "178.571", "14.286"), + "databricks/databricks-gpt-6-luna": ("1.429", "7.143", "1.786", "0.143"), + "databricks/databricks-gpt-6-sol": ("28.571", "142.857", "35.714", "2.857"), "databricks/databricks-claude-opus-5-5": ("57.143", "285.714", "71.429", "2.857"), "databricks/databricks-claude-fable-5-1": ("142.858", "714.286", "178.572", "3.572"), "databricks/databricks-claude-fable-5": ("142.858", "714.286", "178.572", "14.286"), @@ -43,7 +50,7 @@ PUBLISHED_DBU_PER_MILLION: Final = { "databricks/databricks-claude-opus-4-5": ("71.429", "357.143", "89.286", "7.143"), "databricks/databricks-claude-opus-4-1": ("214.286", "1071.429", "267.857", "21.429"), "databricks/databricks-claude-opus-4": ("214.286", "1071.429", "267.857", "21.429"), - "databricks/databricks-claude-sonnet-5": ("42.857", "214.286", "53.571", "4.286"), + "databricks/databricks-claude-sonnet-5": ("28.571", "142.857", "35.714", "2.857"), "databricks/databricks-claude-sonnet-4-6": ("42.857", "214.286", "53.571", "4.286"), "databricks/databricks-claude-sonnet-4-5": ("42.857", "214.286", "53.571", "4.286"), "databricks/databricks-claude-sonnet-4-1": ("42.857", "214.286", "53.571", "4.286"), @@ -197,7 +204,7 @@ def test_models_without_a_cache_discount_bill_cache_tokens_at_the_input_rate( assert prompt_cost > 8000 * info["input_cost_per_token"] -def test_every_model_without_published_cache_dbu_bills_cache_at_its_own_input_rate( +def test_every_model_without_published_cache_dbu_bills_cache_writes_at_its_own_input_rate( local_model_cost_map: None, ) -> None: without_published_rates: Final = [ @@ -210,8 +217,7 @@ def test_every_model_without_published_cache_dbu_bills_cache_at_its_own_input_ra for model in without_published_rates: info = _model_info(model) - for field in CACHE_FIELDS: - assert info[field] == pytest.approx(info["input_cost_per_token"]), (model, field) + assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]), model @pytest.mark.parametrize("model", NEW_MODELS) @@ -224,14 +230,6 @@ def test_backup_price_map_matches_main(model: str) -> None: assert backup_cost[model] == main_cost[model] -def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: None) -> None: - sonnet_5: Final = _model_info("databricks/databricks-claude-sonnet-5") - sonnet_4_6: Final = _model_info("databricks/databricks-claude-sonnet-4-6") - - for field in PRICE_FIELDS: - assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field - - def test_cost_per_token_bills_the_served_priority_tier( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -255,3 +253,29 @@ def test_cost_per_token_bills_the_served_priority_tier( prompt_cost, completion_cost = cost_per_token(model="databricks/dbrx-tiered-test", usage=usage) assert prompt_cost == pytest.approx(30 * 0.001) assert completion_cost == pytest.approx(40 * 0.002) + + +RATES_WITHOUT_A_PRIORITY_FIELD: Final = frozenset({"input_cost_per_image_token"}) + + +def _is_standard_rate_key(key: str) -> bool: + return ( + ("_cost_per_" in key or key.endswith("_token_cost")) + and "dbu" not in key + and key not in RATES_WITHOUT_A_PRIORITY_FIELD + and not key.endswith(("_priority", "_flex", "_batches")) + ) + + +def test_every_databricks_priority_row_prices_every_standard_rate_at_the_priority_tier() -> None: + cost_map: Final = json.loads(MAIN_PRICES.read_text()) + priority_rows: Final = { + key: row + for key, row in cost_map.items() + if key.startswith("databricks/") and "input_cost_per_token_priority" in row + } + + assert priority_rows + for key, row in priority_rows.items(): + missing = sorted(f for f in row if _is_standard_rate_key(f) and f"{f}_priority" not in row) + assert not missing, f"{key} bills these at the standard rate on priority requests: {missing}" diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 04b791cf289..61e5545b8bf 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -1329,7 +1329,7 @@ def test_streaming_surfaces_fireworks_response_fields(): raw_response = MagicMock() raw_response.status_code = 200 raw_response.headers = {} - raw_response.iter_lines = lambda: iter(sse_lines) + raw_response.iter_text = lambda: iter(f"{line}\n" for line in sse_lines) client = HTTPHandler() with patch.object(client, "post", return_value=raw_response): @@ -1788,7 +1788,7 @@ def test_streaming_preserves_selected_model_for_private_accounting(): raw_response = MagicMock() raw_response.status_code = 200 raw_response.headers = {} - raw_response.iter_lines = lambda: iter(sse_lines) + raw_response.iter_text = lambda: iter(f"{line}\n" for line in sse_lines) client = HTTPHandler() with patch.object(client, "post", return_value=raw_response): @@ -2275,7 +2275,7 @@ def test_completion_forwards_litellm_user_id_as_user(call_kwargs: dict[str, obje def test_completion_forwards_litellm_user_id_when_streaming() -> None: client: Final = _fireworks_chat_client() - client.post.return_value.iter_lines = lambda: iter(()) + client.post.return_value.iter_text = lambda: iter(()) litellm.completion( model="fireworks_ai/accounts/fireworks/models/kimi-k3", messages=[{"role": "user", "content": "hi"}], diff --git a/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py index 4119ce99423..53ddb1135b4 100644 --- a/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py +++ b/tests/unit/llms/gemini/google_genai/guardrail_translation/test_google_genai_guardrail_handler.py @@ -4,6 +4,7 @@ Tests for the Google GenAI generateContent guardrail translation handler. import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -190,6 +191,21 @@ async def test_streaming_raw_sse_chunks_accumulate_text_across_split_frames(): assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == ["harmful answer"] +@pytest.mark.asyncio +async def test_streaming_raw_sse_text_holding_unicode_line_separators_reaches_the_guardrail_whole() -> None: + guardrail: Final = _mock_guardrail(["clean"]) + text: Final = "harmful\u2028answer\u2029here" + payload: Final = {"candidates": [{"content": {"parts": [{"text": text}]}}]} + chunk: Final = f"data: {json.dumps(payload, ensure_ascii=False)}\r\n\r\n" + + await GoogleGenAIGenerateContentHandler().process_output_streaming_response( + responses_so_far=[chunk], + guardrail_to_apply=guardrail, + ) + + assert guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] == [text] + + @pytest.mark.asyncio async def test_streaming_blocking_guardrail_exception_propagates(): handler = GoogleGenAIGenerateContentHandler() diff --git a/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py b/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py new file mode 100644 index 00000000000..06485362162 --- /dev/null +++ b/tests/unit/llms/langgraph/chat/test_langgraph_sse_iterator.py @@ -0,0 +1,32 @@ +import json +from typing import Final + +import httpx +import pytest + +from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _messages_event_stream(text: str) -> httpx.Response: + event: Final = ["messages", [[{"type": "AIMessageChunk", "content": text}, {}]]] + return httpx.Response(200, stream=httpx.ByteStream(f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode())) + + +async def _streamed_contents(iterator: LangGraphSSEStreamIterator, sync_stream: bool) -> list[str | None]: + if sync_stream: + return [chunk.choices[0].delta.content for chunk in iterator] + return [chunk.choices[0].delta.content async for chunk in iterator] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_a_message_chunk(sync_stream: bool) -> None: + iterator: Final = LangGraphSSEStreamIterator( + response=_messages_event_stream(_LINE_SEPARATED_TEXT), model="langgraph/agent" + ) + + contents: Final = await _streamed_contents(iterator, sync_stream) + + assert contents == [_LINE_SEPARATED_TEXT, None] diff --git a/tests/unit/llms/oobabooga/chat/test_oobabooga.py b/tests/unit/llms/oobabooga/chat/test_oobabooga.py index 395a4fb5715..332b6913890 100644 --- a/tests/unit/llms/oobabooga/chat/test_oobabooga.py +++ b/tests/unit/llms/oobabooga/chat/test_oobabooga.py @@ -1,7 +1,12 @@ +import json +from typing import Final from unittest.mock import MagicMock, patch +import httpx +import respx import litellm +from litellm.llms.oobabooga.chat import oobabooga MOCK_COMPLETION_RESPONSE = { "choices": [{"message": {"role": "assistant", "content": "hi there"}}], @@ -50,3 +55,27 @@ def test_url_valued_model_still_targets_that_url(): mock_post.assert_called_once() called_url = mock_post.call_args[0][0] assert called_url == "https://sdk-user.example/v1/chat/completions" + + +def test_stream_keeps_unicode_line_separators_inside_a_data_line(respx_mock: respx.MockRouter) -> None: + text: Final = "before\u2028between\u2029after" + chunk: Final = {"choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}]} + data_line: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}" + respx_mock.post("https://ooba.example/v1/chat/completions").mock( + return_value=httpx.Response(200, content=f"{data_line}\n\ndata: [DONE]\n\n".encode()) + ) + + lines: Final = oobabooga.completion( + model="my-model", + messages=[{"role": "user", "content": "hello"}], + api_base="https://ooba.example", + model_response=litellm.ModelResponse(), + print_verbose=print, + encoding=None, + api_key=None, + logging_obj=MagicMock(), + optional_params={"stream": True}, + litellm_params={}, + ) + + assert [line for line in lines if line] == [data_line, "data: [DONE]"] diff --git a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py index 11df07f5fea..e103412e0c3 100644 --- a/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py +++ b/tests/unit/llms/openai/image_generation/test_openai_image_generation_extra_headers.py @@ -265,7 +265,7 @@ class TestImageGenerationEntryPointHeaders: test_headers = {"cf-aig-authorization": "Bearer my-secret"} await litellm.aimage_generation( - model="dall-e-3", + model="gpt-image-2.5-sunburst", prompt="A white cat", extra_headers=test_headers, client=mock_openai_client, diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 831df30c98a..8dc92e23fdf 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -46,12 +46,12 @@ API_FUNCTION_PARAMS = [ ( "image_generation", False, - {"model": "dall-e-3", "prompt": "A beautiful sunset over mountains"}, + {"model": "gpt-image-2.5-sunburst", "prompt": "A beautiful sunset over mountains"}, ), ( "image_generation", True, - {"model": "dall-e-3", "prompt": "A beautiful sunset over mountains"}, + {"model": "gpt-image-2.5-sunburst", "prompt": "A beautiful sunset over mountains"}, ), ( "speech", @@ -84,11 +84,7 @@ async def test_openai_client_reuse(function_name, is_async, args): """ # Determine which client class to mock based on whether the test is async - client_path = ( - "litellm.llms.openai.openai.AsyncOpenAI" - if is_async - else "litellm.llms.openai.openai.OpenAI" - ) + client_path = "litellm.llms.openai.openai.AsyncOpenAI" if is_async else "litellm.llms.openai.openai.OpenAI" # Create the appropriate patches with ( @@ -98,9 +94,7 @@ async def test_openai_client_reuse(function_name, is_async, args): ): # Setup the mock to return None first time (cache miss) then a client for subsequent calls mock_client = MagicMock() - mock_get_cache.side_effect = [None] + [ - mock_client - ] * 9 # First call returns None, rest return the mock client + mock_get_cache.side_effect = [None] + [mock_client] * 9 # First call returns None, rest return the mock client # Make 10 API calls for _ in range(10): @@ -118,9 +112,9 @@ async def test_openai_client_reuse(function_name, is_async, args): pass # Verify client was created only once - assert ( - mock_client_class.call_count == 1 - ), f"{'Async' if is_async else ''}OpenAI client should be created only once" + assert mock_client_class.call_count == 1, ( + f"{'Async' if is_async else ''}OpenAI client should be created only once" + ) # Verify the client was cached assert mock_set_cache.call_count == 1, "Client should be cached once" @@ -144,12 +138,8 @@ def test_precomputed_init_params_match_inspect_signature(): _OPENAI_INIT_PARAMS, ) - expected_openai = tuple( - p for p in inspect.signature(OpenAI.__init__).parameters if p != "self" - ) - expected_azure = tuple( - p for p in inspect.signature(AzureOpenAI.__init__).parameters if p != "self" - ) + expected_openai = tuple(p for p in inspect.signature(OpenAI.__init__).parameters if p != "self") + expected_azure = tuple(p for p in inspect.signature(AzureOpenAI.__init__).parameters if p != "self") assert _OPENAI_INIT_PARAMS == expected_openai assert _AZURE_OPENAI_INIT_PARAMS == expected_azure @@ -500,7 +490,10 @@ async def test_sdk_api_factories_keep_httpx_transport_and_request_timeouts(api: moved: Final = mock.get(host="sdk-default.example", path="/v1/moved-models").mock( return_value=httpx.Response( 200, - json={"object": "list", "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}]}, + json={ + "object": "list", + "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}], + }, ) ) sdk_client: Final = _sdk_api_client(api, is_async, timeout) @@ -542,7 +535,10 @@ async def test_azure_gateway_and_image_variations_use_the_callers_async_session( assert azure_client is not None assert azure_client._client is session response: Final = await aimage_variation( - image=BytesIO(b"image-bytes"), api_key="transport-only", api_base="https://sdk-default.example/v1" + image=BytesIO(b"image-bytes"), + custom_llm_provider="openai", + api_key="transport-only", + api_base="https://sdk-default.example/v1", ) assert response.data[0].url == "https://example.com/image.png" assert not session.is_closed diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py index 75b60b2dd6c..3df1efe311b 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py @@ -4,9 +4,10 @@ from typing import Final from unittest.mock import Mock import httpx +import pytest -from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.llms.openai_like.chat.handler import make_sync_call +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.openai_like.chat.handler import make_call, make_sync_call _FAKE_STREAMED_REPLY: Final = { "id": "chatcmpl-fake-stream", @@ -73,3 +74,35 @@ def test_fake_stream_keeps_reasoning_every_tool_call_and_token_detail(): assert [call.function.name for call in delta.tool_calls] == ["get_weather", "get_time"] assert chunks[0].usage.prompt_tokens_details.cached_tokens == 64 assert chunks[0].usage.completion_tokens_details.reasoning_tokens == 1920 + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _streamed_reply(request: httpx.Request) -> httpx.Response: + chunk: Final = {"id": "chatcmpl-separators", "choices": [{"index": 0, "delta": {"content": _LINE_SEPARATED_TEXT}}]} + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\ndata: [DONE]\n\n" + return httpx.Response(200, stream=httpx.ByteStream(body.encode())) + + +async def _streamed_chunks(sync_stream: bool) -> list[dict[str, object]]: + call_args: Final = { + "api_base": "https://openai-like.example/v1/chat/completions", + "headers": {}, + "data": json.dumps({"model": "openai/gpt-oss-120b", "messages": [], "stream": True}), + "model": "openai/gpt-oss-120b", + "messages": [], + "logging_obj": Mock(), + } + transport: Final = httpx.MockTransport(_streamed_reply) + if sync_stream: + return list(make_sync_call(client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args)) + return [chunk async for chunk in await make_call(client=AsyncHTTPHandler(transport=transport), **call_args)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_the_delta(sync_stream: bool) -> None: + chunks: Final = await _streamed_chunks(sync_stream) + + assert "".join(str(chunk["text"]) for chunk in chunks) == _LINE_SEPARATED_TEXT diff --git a/tests/unit/llms/openai_like/test_zerogpu_provider.py b/tests/unit/llms/openai_like/test_zerogpu_provider.py new file mode 100644 index 00000000000..1021237c335 --- /dev/null +++ b/tests/unit/llms/openai_like/test_zerogpu_provider.py @@ -0,0 +1,239 @@ +import json +from collections.abc import Mapping +from typing import Final + +import httpx +import pytest +import respx +from openai import OpenAI + +import litellm +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.proxy.auth.model_checks import get_provider_models + +ZEROGPU_API_BASE: Final = "https://api.zerogpu.ai/v1" + +ZEROGPU_MODELS: Final = ( + "deepseek-v4.1-flash", + "glm-5.3-flash", + "gpt-4.1-mini", + "gpt-5.4-nano", + "gpt-5.6-luna", + "gpt-6-luna", + "gpt-oss-120b", + "LFM2.5-1.2B-Instruct", + "LFM2.5-1.2B-Thinking", + "llama-3.1-8b-instruct-fast", + "llama-3.2-3b-instruct-fast", + "qwen3-30b-a3b-fp8", + "qwen3.6-35b-a3b", +) + + +WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, +} + + +class _RecordingZeroGPU: + def __init__(self, usage: Mapping[str, object]) -> None: + self.usage: Final = usage + self.requests: tuple[httpx.Request, ...] = () + + def respond(self, request: httpx.Request) -> httpx.Response: + self.requests = (*self.requests, request) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-zerogpu", + "object": "chat.completion", + "created": 1234567890, + "model": json.loads(request.content)["model"], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": dict(self.usage), + }, + ) + + def client(self) -> OpenAI: + return OpenAI( + api_key="zerogpu-test-key", + base_url=ZEROGPU_API_BASE, + http_client=httpx.Client(transport=httpx.MockTransport(self.respond)), + ) + + +def test_prefixed_model_reaches_zerogpu_chat_completions_without_the_prefix(): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + litellm.completion( + model="zerogpu/gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + assert len(zerogpu.requests) == 1 + assert str(zerogpu.requests[0].url) == f"{ZEROGPU_API_BASE}/chat/completions" + assert json.loads(zerogpu.requests[0].content) == { + "model": "gpt-oss-120b", + "messages": [{"role": "user", "content": "hello"}], + } + + +def test_responses_request_reaches_zerogpu_responses_natively(): + with respx.mock() as upstream: + route: Final = upstream.post(f"{ZEROGPU_API_BASE}/responses").respond( + 200, + json={ + "id": "resp_zerogpu", + "object": "response", + "created_at": 1_792_000_000, + "model": "gpt-oss-120b", + "status": "completed", + "output": [ + { + "id": "msg_zerogpu", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4}, + }, + ) + response: Final = litellm.responses(model="zerogpu/gpt-oss-120b", input="hello", api_key="zerogpu-test-key") + + body: Final = json.loads(route.calls.last.request.content) + assert body["model"] == "gpt-oss-120b" + assert body["input"] == "hello" + assert response.output[0].content[0].text == "hi" + + +def test_anthropic_messages_request_is_bridged_to_zerogpu_chat_completions(): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + response: Final = litellm.anthropic.messages.create( + model="zerogpu/gpt-oss-120b", + messages=[{"role": "user", "content": "hello"}], + max_tokens=10, + client=client, + ) + + assert [str(request.url) for request in zerogpu.requests] == [f"{ZEROGPU_API_BASE}/chat/completions"] + assert json.loads(zerogpu.requests[0].content) == { + "model": "gpt-oss-120b", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + } + assert response["content"] == [{"type": "text", "text": "hi"}] + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_zerogpu_completion_is_charged_at_the_cost_map_rates(model: str): + usage: Final = { + "prompt_tokens": 1000, + "completion_tokens": 200, + "total_tokens": 1200, + "prompt_tokens_details": {"cached_tokens": 400}, + } + zerogpu: Final = _RecordingZeroGPU(usage) + + with zerogpu.client() as client: + response: Final = litellm.completion( + model=f"zerogpu/{model}", + messages=[{"role": "user", "content": "hello"}], + client=client, + ) + + rates: Final = litellm.model_cost[f"zerogpu/{model}"] + cached_rate: Final = rates.get("cache_read_input_token_cost", rates["input_cost_per_token"]) + expected_cost: Final = ( + 600 * rates["input_cost_per_token"] + 400 * cached_rate + 200 * rates["output_cost_per_token"] + ) + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) + assert expected_cost > 0 + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_trim_messages_fits_zerogpu_history_into_the_model_context_window(model: str): + budget: Final = 20 + trim_ratio: Final = budget / litellm.model_cost[f"zerogpu/{model}"]["max_input_tokens"] + messages: Final = [{"role": "user", "content": "lorem ipsum " * 300}] + + trimmed: Final = litellm.utils.trim_messages(messages, model=f"zerogpu/{model}", trim_ratio=trim_ratio) + + trimmed_tokens: Final = litellm.token_counter(model=f"zerogpu/{model}", messages=trimmed) + assert trimmed_tokens <= budget, trimmed + + +@pytest.mark.parametrize( + ("model", "forwards_tools", "forwards_reasoning_effort"), + [ + ("gpt-5.6-luna", True, True), + ("gpt-oss-120b", True, True), + ("qwen3.6-35b-a3b", True, True), + ("LFM2.5-1.2B-Instruct", False, False), + ], +) +def test_tools_and_reasoning_effort_reach_zerogpu_only_for_models_that_support_them( + model: str, forwards_tools: bool, forwards_reasoning_effort: bool +): + zerogpu: Final = _RecordingZeroGPU({"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}) + + with zerogpu.client() as client: + litellm.completion( + model=f"zerogpu/{model}", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=[WEATHER_TOOL], + tool_choice="auto", + reasoning_effort="high", + drop_params=True, + client=client, + ) + + tools: Final = {"tools": [WEATHER_TOOL], "tool_choice": "auto"} if forwards_tools else {} + reasoning: Final = {"reasoning_effort": "high"} if forwards_reasoning_effort else {} + assert json.loads(zerogpu.requests[0].content) == { + "model": model, + "messages": [{"role": "user", "content": "weather in Paris?"}], + **tools, + **reasoning, + } + + +@pytest.mark.parametrize("model", ZEROGPU_MODELS) +def test_model_info_reports_the_reasoning_support_litellm_routes_with(model: str): + model_info: Final = litellm.get_model_info(f"zerogpu/{model}") + + assert bool(model_info.get("supports_reasoning")) == litellm.utils.supports_reasoning(f"zerogpu/{model}") + + +def test_zerogpu_key_and_base_come_from_zerogpu_env_vars(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ZEROGPU_API_KEY", "zerogpu-env-key") + monkeypatch.delenv("ZEROGPU_API_BASE", raising=False) + + assert get_llm_provider(model="zerogpu/gpt-oss-120b") == ( + "gpt-oss-120b", + "zerogpu", + "zerogpu-env-key", + ZEROGPU_API_BASE, + ) + + monkeypatch.setenv("ZEROGPU_API_BASE", "https://zerogpu.internal.example/v1") + + assert get_llm_provider(model="zerogpu/gpt-oss-120b")[3] == "https://zerogpu.internal.example/v1" + + +def test_zerogpu_wildcard_lists_every_zerogpu_catalog_model(): + catalog_models: Final = sorted( + key for key, row in litellm.model_cost.items() if row.get("litellm_provider") == "zerogpu" + ) + + assert catalog_models + assert sorted(get_provider_models("zerogpu") or []) == catalog_models diff --git a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py index f9d12b7eeea..36c319d9524 100644 --- a/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py +++ b/tests/unit/llms/snowflake/chat/test_snowflake_chat_transformation.py @@ -902,20 +902,14 @@ class TestSnowflakeChatCompletion: @pytest.mark.parametrize("sync_mode", [True, False]) def test_chat_completion_snowflake_stream(self, sync_mode): - raw_chunks = _streaming_chunks() + sse_body = "".join(f"data: {chunk}\n\n" for chunk in _streaming_chunks()) + "data: [DONE]\n\n" + mock_resp = httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=httpx.ByteStream(sse_body.encode()), + ) if sync_mode: - - def _iter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.iter_lines.return_value = _iter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - with patch.object(HTTPHandler, "post", return_value=mock_resp) as mock_post: response = completion( model="snowflake/mistral-7b", @@ -930,16 +924,6 @@ class TestSnowflakeChatCompletion: mock_post.assert_called_once() else: - async def _aiter_lines(): - for chunk in raw_chunks: - yield f"data: {chunk}" - yield "data: [DONE]" - - mock_resp = MagicMock() - mock_resp.aiter_lines.return_value = _aiter_lines() - mock_resp.status_code = 200 - mock_resp.headers = {"content-type": "text/event-stream"} - async def _run(): with patch.object( AsyncHTTPHandler, @@ -964,8 +948,8 @@ class TestSnowflakeChatCompletion: chunks_received = asyncio.run(_run()) - assert len(chunks_received) > 0 content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content) + assert content == "The sky is blue" async def test_snowflake_claude_async_completion_inlines_remote_images_off_the_event_loop(async_only_image_fetch): diff --git a/tests/unit/llms/tencent/image_edit/__init__.py b/tests/unit/llms/tencent/image_edit/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/image_edit/test_tencent_image_edit_transformation.py b/tests/unit/llms/tencent/image_edit/test_tencent_image_edit_transformation.py new file mode 100644 index 00000000000..506c43f6cba --- /dev/null +++ b/tests/unit/llms/tencent/image_edit/test_tencent_image_edit_transformation.py @@ -0,0 +1,232 @@ +import base64 +import io +import json +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +PNG_BYTES: Final = b"\x89PNG\r\n\x1a\n" + b"\x00" * 32 +TOKENHUB_RESPONSE: Final[dict[str, object]] = { + "created": 1792551838, + "data": [{"output_format": "jpeg", "size": "2048x2048", "url": "https://cos.example/edited.jpg"}], + "model": "doubao-seedream-5-0-lite", + "tokenhub_usage": {"total_tokens": 21875}, + "request_id": "4f795d48-b227-4c35-af2c-1a0ba288b6b6", +} + + +def _client(captured: dict[str, object]) -> HTTPHandler: + def respond(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["content_type"] = request.headers.get("content-type") + captured["authorization"] = request.headers.get("authorization") + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=TOKENHUB_RESPONSE) + + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + + +def test_edit_sends_uploaded_images_as_data_urls_in_a_json_body() -> None: + captured: Final[dict[str, object]] = {} + png_data_url: Final = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}" + + response: Final = litellm.image_edit( + model="tencent/seedream-image-v5.0-lite", + image=[PNG_BYTES, io.BytesIO(PNG_BYTES), "https://ref.example/style.png"], + prompt="turn the square into a paper lantern", + size="2K", + api_key="sk-tokenhub", + client=_client(captured), + ) + + assert captured["url"] == "https://tokenhub-intl.tencentcloudmaas.com/v1/wand/si-image/generation" + assert captured["content_type"] == "application/json" + assert captured["authorization"] == "Bearer sk-tokenhub" + assert captured["body"] == { + "model": "seedream-image-v5.0-lite", + "prompt": "turn the square into a paper lantern", + "images": [png_data_url, png_data_url, "https://ref.example/style.png"], + "size": "2K", + } + assert response.data[0].url == "https://cos.example/edited.jpg" + assert response.usage.total_tokens == 21875 + + +def test_edit_rejects_a_mask_without_drop_params() -> None: + captured: Final[dict[str, object]] = {} + + with pytest.raises(litellm.UnsupportedParamsError): + litellm.image_edit( + model="tencent/seedream-image-v5.0-pro", + image=PNG_BYTES, + mask=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + ) + + assert captured == {} + + +def test_edit_drops_a_mask_under_drop_params() -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-pro", + image=PNG_BYTES, + mask=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + drop_params=True, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-pro", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + } + + +def test_edit_is_not_offered_for_tencent_models_outside_seedream() -> None: + captured: Final[dict[str, object]] = {} + + with pytest.raises(Exception, match="image edit is not supported for tencent"): + litellm.image_edit( + model="tencent/hy-image-v3.5-preview", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + ) + + assert captured == {} + + +@pytest.mark.parametrize( + ("form_params", "expected_body_params"), + [ + ({"watermark": "false", "seed": "7"}, {"watermark": False, "seed": 7}), + ({"watermark": "True", "seed": "-1"}, {"watermark": True, "seed": -1}), + ({"watermark": "maybe"}, {"watermark": "maybe"}), + ( + {"layer_decomposition": "true", "optimize_prompt_options[mode]": "fast"}, + {"layer_decomposition": True, "optimize_prompt_options": {"mode": "fast"}}, + ), + ( + {"sequential_image_generation_options[max_images]": "3", "user": "12345"}, + {"sequential_image_generation_options": {"max_images": 3}, "user": "12345"}, + ), + ], +) +def test_edit_parses_tokenhub_values_sent_as_form_strings( + form_params: dict[str, object], expected_body_params: dict[str, object] +) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-flash", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + **form_params, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-flash", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + **expected_body_params, + } + + +@pytest.mark.parametrize( + ("params", "expected_body_params"), + [ + ({"background": "auto", "watermark": False}, {"watermark": False}), + ({"watermark": False, "seed": 7, "additional_drop_params": ["watermark"]}, {"seed": 7}), + ({"extra_body": {"watermark": False}, "additional_drop_params": ["watermark"]}, {}), + ], +) +def test_edit_keeps_dropped_params_out_of_the_tokenhub_body( + params: dict[str, object], expected_body_params: dict[str, object] +) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-flash", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + **params, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-flash", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + **expected_body_params, + } + + +@pytest.mark.parametrize( + ("params", "expected_body_params"), + [ + ({"extra_body": {"response_format": "b64_json"}}, {"response_format": "b64_json"}), + ({"size": "1K", "extra_body": {"size": "2K"}}, {"size": "2K"}), + ({"extra_body": {"model": "seedream-image-v5.0-pro", "seed": 7}}, {"seed": 7}), + ({"extra_body": '{"seed": 7}'}, {}), + ], +) +def test_edit_sends_extra_body_fields_as_given( + params: dict[str, object], expected_body_params: dict[str, object] +) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-flash", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + **params, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-flash", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + **expected_body_params, + } + + +def test_edit_forwards_tokenhub_only_params_and_extra_body() -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_edit( + model="tencent/seedream-image-v5.0-flash", + image=PNG_BYTES, + prompt="turn the square into a paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + user="end-user-42", + watermark=False, + seed=7, + extra_body={"sequential_image_generation": "disabled"}, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-flash", + "prompt": "turn the square into a paper lantern", + "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"], + "user": "end-user-42", + "watermark": False, + "seed": 7, + "sequential_image_generation": "disabled", + } diff --git a/tests/unit/llms/tencent/image_generation/__init__.py b/tests/unit/llms/tencent/image_generation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/tencent/image_generation/test_tencent_image_generation_transformation.py b/tests/unit/llms/tencent/image_generation/test_tencent_image_generation_transformation.py new file mode 100644 index 00000000000..ab16382cdae --- /dev/null +++ b/tests/unit/llms/tencent/image_generation/test_tencent_image_generation_transformation.py @@ -0,0 +1,215 @@ +import json +from typing import Final + +import httpx +import openai +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +TOKENHUB_RESPONSE: Final[dict[str, object]] = { + "created": 1792551825, + "data": [{"output_format": "jpeg", "size": "1248x832", "url": "https://cos.example/lantern.jpg"}], + "model": "doubao-seedream-5-0-pro-260628", + "tokenhub_usage": {"total_tokens": 28125}, + "request_id": "c57b0374-b5fe-47db-b4b2-751315f920aa", +} + + +def _client(captured: dict[str, object], status_code: int = 200, body: dict[str, object] | None = None) -> HTTPHandler: + def respond(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["authorization"] = request.headers.get("authorization") + captured["body"] = json.loads(request.content) + return httpx.Response(status_code, json=TOKENHUB_RESPONSE if body is None else body) + + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + + +def test_generation_posts_openai_and_tokenhub_params_to_the_wand_route() -> None: + captured: Final[dict[str, object]] = {} + + response: Final = litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + size="2K", + response_format="url", + api_key="sk-tokenhub", + client=_client(captured), + images=["https://ref.example/lantern.png"], + watermark=False, + ) + + assert captured["url"] == "https://tokenhub-intl.tencentcloudmaas.com/v1/wand/si-image/generation" + assert captured["authorization"] == "Bearer sk-tokenhub" + assert captured["body"] == { + "model": "seedream-image-v5.0-pro", + "prompt": "a red paper lantern", + "size": "2K", + "response_format": "url", + "images": ["https://ref.example/lantern.png"], + "watermark": False, + } + assert response.data[0].url == "https://cos.example/lantern.jpg" + assert response.size == "1248x832" + assert response.output_format == "jpeg" + assert response.created == 1792551825 + assert response.usage.total_tokens == 28125 + assert response.usage.output_tokens == 28125 + + +def test_generation_returns_base64_images() -> None: + captured: Final[dict[str, object]] = {} + body: Final[dict[str, object]] = { + **TOKENHUB_RESPONSE, + "data": [{"output_format": "png", "size": "1024x1024", "b64_json": "aW1n"}], + } + + response: Final = litellm.image_generation( + model="tencent/seedream-image-v5.0-flash", + prompt="a red paper lantern", + response_format="b64_json", + api_key="sk-tokenhub", + client=_client(captured, body=body), + ) + + assert response.data[0].b64_json == "aW1n" + assert response.data[0].url is None + + +@pytest.mark.parametrize( + ("api_base", "expected_url"), + [ + ("https://tokenhub.example", "https://tokenhub.example/v1/wand/si-image/generation"), + ("https://tokenhub.example/v1", "https://tokenhub.example/v1/wand/si-image/generation"), + ("https://tokenhub.example/v1/", "https://tokenhub.example/v1/wand/si-image/generation"), + ], +) +def test_generation_resolves_a_configured_api_base_to_the_wand_route(api_base: str, expected_url: str) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_generation( + model="tencent/seedream-image-v5.0-lite", + prompt="a red paper lantern", + api_key="sk-tokenhub", + api_base=api_base, + client=_client(captured), + ) + + assert captured["url"] == expected_url + + +@pytest.mark.parametrize( + ("params", "expected_body_params"), + [ + ({"n": 1}, {}), + ({"n": 3, "drop_params": True}, {}), + ({"background": "auto"}, {}), + ({"background": "transparent", "output_format": "png"}, {"background": "transparent", "output_format": "png"}), + ({"user": "end-user-42"}, {"user": "end-user-42"}), + ], +) +def test_generation_sends_only_values_tokenhub_accepts( + params: dict[str, object], expected_body_params: dict[str, object] +) -> None: + captured: Final[dict[str, object]] = {} + + litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + api_key="sk-tokenhub", + client=_client(captured), + **params, + ) + + assert captured["body"] == { + "model": "seedream-image-v5.0-pro", + "prompt": "a red paper lantern", + **expected_body_params, + } + + +def test_generation_rejects_more_than_one_image_without_drop_params() -> None: + captured: Final[dict[str, object]] = {} + + with pytest.raises(litellm.UnsupportedParamsError) as exc_info: + litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + n=2, + api_key="sk-tokenhub", + client=_client(captured), + ) + + assert exc_info.value.status_code == 400 + assert captured == {} + + +@pytest.mark.parametrize( + ("error", "expected_message"), + [ + ( + { + "error": { + "message": "The model or service ID 'seedream-image-v5.0-pro' does not exist.", + "code": "400004", + "type": "gateway_error", + } + }, + "does not exist", + ), + ({"Type": 2, "Code": 1001, "Msg": "Invalid param", "Desc": ""}, "Invalid param"), + ], +) +def test_generation_surfaces_the_tokenhub_error_message(error: dict[str, object], expected_message: str) -> None: + with pytest.raises(litellm.BadRequestError) as exc_info: + litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + api_key="sk-tokenhub", + client=_client({}, status_code=400, body=error), + ) + + assert expected_message in str(exc_info.value) + + +def test_generation_without_a_tokenhub_key_is_an_authentication_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("TENCENT_API_KEY", raising=False) + captured: Final[dict[str, object]] = {} + + with pytest.raises(litellm.AuthenticationError) as exc_info: + litellm.image_generation( + model="tencent/seedream-image-v5.0-pro", + prompt="a red paper lantern", + client=_client(captured), + ) + + assert exc_info.value.status_code == 401 + assert "TENCENT_API_KEY" in str(exc_info.value) + assert captured == {} + + +def test_generation_keeps_other_tencent_image_models_on_the_openai_compatible_route() -> None: + captured: Final[dict[str, object]] = {} + + def respond(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["body"] = json.loads(request.content) + return httpx.Response(200, json={"created": 1792551825, "data": [{"url": "https://cos.example/hy.png"}]}) + + response: Final = litellm.image_generation( + model="tencent/hy-image-v3.5-preview", + prompt="a red paper lantern", + quality="high", + api_key="sk-tokenhub", + client=openai.OpenAI( + api_key="sk-tokenhub", + base_url="https://tokenhub-intl.tencentcloudmaas.com/v1", + http_client=httpx.Client(transport=httpx.MockTransport(respond)), + ), + ) + + assert captured["url"] == "https://tokenhub-intl.tencentcloudmaas.com/v1/images/generations" + assert captured["body"] == {"model": "hy-image-v3.5-preview", "prompt": "a red paper lantern", "quality": "high"} + assert response.data[0].url == "https://cos.example/hy.png" diff --git a/tests/unit/llms/tencent/test_tencent_cost_calculator.py b/tests/unit/llms/tencent/test_tencent_cost_calculator.py new file mode 100644 index 00000000000..7f98ee02810 --- /dev/null +++ b/tests/unit/llms/tencent/test_tencent_cost_calculator.py @@ -0,0 +1,128 @@ +from typing import Final + +import httpx +import openai +import pytest + +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails + +SEEDREAM_MODELS: Final = ( + "tencent/seedream-image-v5.0-pro", + "tencent/seedream-image-v5.0-flash", + "tencent/seedream-image-v5.0-lite", +) + + +def _image_response(total_tokens: int, images: int = 1) -> ImageResponse: + return ImageResponse( + created=1792551825, + data=[ImageObject(url=f"https://cos.example/{index}.jpg") for index in range(images)], + usage=ImageUsage( + input_tokens=0, + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0), + output_tokens=total_tokens, + total_tokens=total_tokens, + ), + ) + + +@pytest.mark.parametrize("model", SEEDREAM_MODELS) +@pytest.mark.parametrize("call_type", ["image_generation", "image_edit"]) +def test_seedream_bills_tokenhub_usage_tokens_for_a_tier_size( + local_model_cost_map: None, model: str, call_type: str +) -> None: + token_rate: Final = litellm.model_cost[model]["output_cost_per_image_token"] + + cost: Final = litellm.completion_cost( + completion_response=_image_response(total_tokens=56250), + model=model, + call_type=call_type, + size="2K", + ) + + assert cost == pytest.approx(56250 * token_rate) + + +@pytest.mark.parametrize("model", SEEDREAM_MODELS) +def test_seedream_falls_back_to_the_per_image_price_without_usage(local_model_cost_map: None, model: str) -> None: + per_image_price: Final = litellm.model_cost[model]["output_cost_per_image"] + + cost: Final = litellm.completion_cost( + completion_response=_image_response(total_tokens=0, images=2), + model=model, + call_type="image_generation", + size="2K", + ) + + assert cost == pytest.approx(2 * per_image_price) + + +async def test_a_deployment_priced_tencent_image_model_outside_seedream_bills_its_deployment_price() -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"created": 1792551825, "data": [{"url": "https://cos.example/hy.png"}]}) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "hy-image", + "litellm_params": {"model": "tencent/hy-image-v3.5-preview", "api_key": "sk-tokenhub"}, + "model_info": {"input_cost_per_image": 0.04}, + } + ] + ) + + response: Final = await router.aimage_generation( + model="hy-image", + prompt="a red paper lantern", + client=openai.AsyncOpenAI( + api_key="sk-tokenhub", + base_url="https://tokenhub-intl.tencentcloudmaas.com/v1", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ), + ) + + assert response._hidden_params["response_cost"] == pytest.approx(0.04) + + +@pytest.mark.parametrize( + ("deployment_price", "expected_cost"), + [ + ({"output_cost_per_token": 0.000001}, 11250 * 0.000001), + ({"output_cost_per_image": 0.05}, 0.05), + ], +) +def test_a_seedream_deployment_priced_in_the_ui_bills_its_own_rate( + deployment_price: dict[str, float], expected_cost: float +) -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "created": 1792551825, + "data": [{"output_format": "jpeg", "size": "1024x1024", "url": "https://cos.example/0.jpg"}], + "tokenhub_usage": {"total_tokens": 11250}, + }, + ) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "seedream-flash", + "litellm_params": { + "model": "tencent/seedream-image-v5.0-flash", + "api_key": "sk-tokenhub", + **deployment_price, + }, + } + ] + ) + + response: Final = router.image_generation( + model="seedream-flash", + prompt="a red paper lantern", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + ) + + assert response._hidden_params["response_cost"] == pytest.approx(expected_cost) diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py index a02840ed4ef..7afbd05b221 100644 --- a/tests/unit/llms/test_predibase_transformation.py +++ b/tests/unit/llms/test_predibase_transformation.py @@ -1,3 +1,4 @@ +import json from typing import Final from unittest.mock import AsyncMock, Mock @@ -639,3 +640,49 @@ def test_predibase_completion_passes_existing_config_to_async_completion(monkeyp assert result == "async-result" assert captured["async_kwargs"]["predibase_config"] is captured["config_instance"] + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _generate_stream_body(text: str) -> bytes: + events: Final = ( + {"token": {"text": text}}, + {"token": {"text": ""}, "generated_text": text, "details": {"finish_reason": "eos_token"}}, + ) + return "".join(f"data:{json.dumps(event, ensure_ascii=False)}\n\n" for event in events).encode() + + +async def _streamed_contents(sync_stream: bool) -> list[str | None]: + request: Final = { + "model": "predibase/llama-3-8b-instruct", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "predibase_tenant_id": "tenant-1", + "api_key": "pb-key", + "max_retries": 0, + } + if sync_stream: + return [chunk.choices[0].delta.content for chunk in litellm.completion(**request)] + return [chunk.choices[0].delta.content async for chunk in await litellm.acompletion(**request)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_predibase_stream_keeps_unicode_line_separators_inside_a_token( + sync_stream: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_hf_tokenizer_download", True) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr( + litellm, + "custom_prompt_dict", + {"llama-3-8b-instruct": {"roles": {}, "initial_prompt_value": "", "final_prompt_value": ""}}, + ) + respx_mock.post( + "https://serving.app.predibase.com/tenant-1/deployments/v2/llms/llama-3-8b-instruct/generate_stream" + ).mock(return_value=httpx.Response(200, content=_generate_stream_body(_LINE_SEPARATED_TEXT))) + + contents: Final = await _streamed_contents(sync_stream) + + assert contents == [_LINE_SEPARATED_TEXT, None] diff --git a/tests/unit/llms/triton/test_triton.py b/tests/unit/llms/triton/test_triton.py index 7c0de603a3e..238c30816b2 100644 --- a/tests/unit/llms/triton/test_triton.py +++ b/tests/unit/llms/triton/test_triton.py @@ -1,5 +1,7 @@ import json import traceback +from collections.abc import Iterator +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -210,8 +212,8 @@ def test_completion_triton_generate_api(stream, monkeypatch, respx_mock): mock_response = MagicMock() if stream: - def mock_iter_lines(): - mock_output = "".join( + def mock_iter_text() -> Iterator[str]: + mock_output: Final = "".join( [ 'data: {"model_name":"ensemble","model_version":"1","sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"' + t @@ -219,10 +221,9 @@ def test_completion_triton_generate_api(stream, monkeypatch, respx_mock): for t in ["I", " am", " an", " AI", " assistant"] ] ) - for out in mock_output.split("\n"): - yield out + yield mock_output - mock_response.iter_lines = mock_iter_lines + mock_response.iter_text = mock_iter_text else: def return_val(): diff --git a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py index 06aaca44e8a..20333fd0465 100644 --- a/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py +++ b/tests/unit/llms/vertex_ai/agent_engine/test_transformation.py @@ -14,7 +14,7 @@ import pytest from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, get_async_httpx_client from litellm.llms.vertex_ai.agent_engine.sse_iterator import ( VertexAgentEngineResponseIterator, ) @@ -170,3 +170,48 @@ async def test_async_stream_wrapper_without_a_client_posts_through_the_cached_ve assert type(wrapper) is CustomStreamWrapper assert [str(request.url) for request in sent_requests] == [api_base] assert json.loads(sent_requests[0].content) == {"class_method": "stream_query"} + + +_LINE_SEPARATED_TEXT: Final = "before\u2028between\u2029after" + + +def _separated_text_reply(request: httpx.Request) -> httpx.Response: + chunk: Final = {"content": {"parts": [{"text": _LINE_SEPARATED_TEXT}], "role": "model"}, "finish_reason": "STOP"} + return httpx.Response(200, stream=httpx.ByteStream(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode())) + + +async def _streamed_text(sync_stream: bool) -> str: + messages: Final = [{"role": "user", "content": "hi"}] + call_args: Final = { + "model": "agent_engine/123", + "custom_llm_provider": "vertex_ai", + "logging_obj": Logging( + model="agent_engine/123", + messages=messages, + stream=True, + call_type="completion" if sync_stream else "acompletion", + start_time=datetime(2026, 1, 1), + litellm_call_id="call-separators", + function_id="fn-separators", + ), + "api_base": "https://us-central1-aiplatform.googleapis.com/v1/reasoningEngines/123:streamQuery", + "headers": {}, + "data": {"class_method": "stream_query"}, + "messages": messages, + "litellm_params": {}, + } + transport: Final = httpx.MockTransport(_separated_text_reply) + if sync_stream: + sync_wrapper: Final = VertexAgentEngineConfig().get_sync_custom_stream_wrapper( + client=HTTPHandler(client=httpx.Client(transport=transport)), **call_args + ) + return "".join(chunk.choices[0].delta.content or "" for chunk in sync_wrapper) + async_wrapper: Final = await VertexAgentEngineConfig().get_async_custom_stream_wrapper( + client=AsyncHTTPHandler(transport=transport), **call_args + ) + return "".join([chunk.choices[0].delta.content or "" async for chunk in async_wrapper]) + + +@pytest.mark.parametrize("sync_stream", [True, False]) +async def test_stream_keeps_unicode_line_separators_inside_the_text(sync_stream: bool) -> None: + assert await _streamed_text(sync_stream) == _LINE_SEPARATED_TEXT diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 065d894e658..67c78b62151 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -3,8 +3,9 @@ import json import logging import re import threading -from collections.abc import Awaitable, Callable, Iterator, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping from copy import deepcopy +from itertools import pairwise from typing import Final, List, Literal, cast, get_args from unittest.mock import MagicMock, Mock, patch @@ -34,7 +35,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ) from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata -from litellm.types.utils import ChoiceLogprobs, ModelInfo, Usage +from litellm.types.utils import ChoiceLogprobs, ModelInfo, ModelResponseStream, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.integrations.custom_logger import CustomLogger @@ -3514,6 +3515,39 @@ def test_vertex_ai_annotation_conversion(): assert "Current time" in time_annotation["url_citation"]["title"] +@pytest.mark.parametrize( + "indices, expected_indices", + [([2, 0], [2, 0]), ([1, 99, -1, 2, 0], [2, 0]), ([0], [0]), ([], []), ([1, 99], [])], + ids=["multiple-sources", "non-web-first", "single-source", "empty", "no-web"], +) +def test_grounding_annotations_preserve_all_web_sources(indices: list[int], expected_indices: list[int]) -> None: + chunks: Final = [ + {"web": {"uri": "https://example.com/first", "title": "First"}}, + {"retrievedContext": {"uri": "gs://documents/source"}}, + {"web": {"uri": "https://example.com/second", "title": "Second"}}, + ] + metadata: Final = { + "groundingChunks": chunks, + "groundingSupports": [ + {"segment": {"startIndex": 4, "endIndex": 11}, "groundingChunkIndices": indices} + ], + } + annotations: Final = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata], "The passage") + expected: Final = [ + { + "type": "url_citation", + "url_citation": { + "start_index": 4, + "end_index": 11, + "url": chunks[index]["web"]["uri"], + "title": chunks[index]["web"]["title"], + }, + } + for index in expected_indices + ] + assert json.dumps(annotations) == json.dumps(expected) + + def test_vertex_ai_annotation_empty_grounding_metadata(): """ Test handling of empty or missing grounding metadata. @@ -10092,7 +10126,7 @@ def test_vertex_ai_streaming_response_id(): mock_response = MagicMock() mock_response.status_code = 200 - mock_response.iter_lines = MagicMock(return_value=stream_response()) + mock_response.iter_text = MagicMock(return_value=stream_response()) return mock_response logging_obj = MagicMock() @@ -10655,3 +10689,91 @@ async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_n sent: Final = await _sent_thinking_config("gemini-3.8-flash", {"reasoning_effort": "high"}) assert sent["includeThoughts"] is expected + + +class _BodyInChunks(httpx.SyncByteStream, httpx.AsyncByteStream): + def __init__(self, chunks: tuple[bytes, ...]) -> None: + self._chunks: Final = chunks + + def __iter__(self) -> Iterator[bytes]: + yield from self._chunks + + async def __aiter__(self) -> AsyncIterator[bytes]: + for chunk in self._chunks: + yield chunk + + +def _split_at(body: bytes, cuts: tuple[int, ...]) -> tuple[bytes, ...]: + return tuple(body[start:end] for start, end in pairwise((0, *sorted(cuts), len(body)))) + + +async def _gemini_stream(is_async: bool, chunks: tuple[bytes, ...]) -> list[ModelResponseStream]: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, stream=_BodyInChunks(chunks), headers={"content-type": "text/event-stream"}) + + kwargs: Final = { + "model": "gemini/gemini-3.8-flash", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + "stream_options": {"include_usage": True}, + "api_key": "fake-gemini-key", + } + if is_async: + stream: Final = await litellm.acompletion( + **kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(handler)) + ) + return [chunk async for chunk in stream] + return list( + litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))) + ) + + +@pytest.mark.parametrize("is_async", [True, False]) +async def test_gemini_stream_keeps_text_and_usage_when_a_part_contains_a_unicode_line_separator(is_async: bool) -> None: + events: Final = ( + {"candidates": [{"content": {"role": "model", "parts": [{"text": "line one\u2028line two"}]}}]}, + { + "candidates": [{"content": {"role": "model", "parts": [{"text": " and three"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 8, "totalTokenCount": 18}, + }, + ) + body: Final = "".join(f"data: {json.dumps(event, ensure_ascii=False)}\r\n\r\n" for event in events).encode() + cut_inside_separator_and_crlf: Final = (body.index(b"\xe2\x80\xa8") + 1, body.index(b"\r") + 1) + + received: Final = await _gemini_stream(is_async, _split_at(body, cut_inside_separator_and_crlf)) + + usage: Final = next(chunk.usage for chunk in reversed(received) if getattr(chunk, "usage", None)) + text: Final = "".join(c.choices[0].delta.content or "" for c in received if c.choices) + assert text == "line one\u2028line two and three" + assert (usage.prompt_tokens, usage.completion_tokens) == (10, 8) + + +@pytest.mark.parametrize("chunk_indices", [[0, 1], [2, 1, 0]]) +def test_grounding_support_preserves_every_web_source(chunk_indices: list[int]) -> None: + metadata: Final = { + "groundingChunks": [ + {"web": {"uri": "https://example.com/a", "title": "A"}}, + {"web": {"uri": "https://example.com/b", "title": "B"}}, + {"retrievedContext": {"uri": "gs://document"}}, + ], + "groundingSupports": [ + {"segment": {"startIndex": 0, "endIndex": 5}, "groundingChunkIndices": chunk_indices}, + {"segment": {"startIndex": 6, "endIndex": 11}, "groundingChunkIndices": [1]}, + ], + } + + annotations: Final = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata], "First claim") + + expected_sources: Final = [(index, 0, 5) for index in chunk_indices if index in (0, 1)] + [(1, 6, 11)] + assert annotations == [ + { + "type": "url_citation", + "url_citation": { + "start_index": start, + "end_index": end, + "url": metadata["groundingChunks"][index]["web"]["uri"], + "title": metadata["groundingChunks"][index]["web"]["title"], + }, + } + for index, start, end in expected_sources + ] diff --git a/tests/unit/llms/vertex_ai/test_http_status_201.py b/tests/unit/llms/vertex_ai/test_http_status_201.py index 34502d5cebe..0ed469acfa7 100644 --- a/tests/unit/llms/vertex_ai/test_http_status_201.py +++ b/tests/unit/llms/vertex_ai/test_http_status_201.py @@ -39,8 +39,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): # Create a mock response with status code 201 mock_response = MagicMock() mock_response.status_code = 201 - mock_response.aiter_lines = MagicMock() - mock_response.aiter_lines.return_value = ["test response"] + mock_response.aiter_text = MagicMock() + mock_response.aiter_text.return_value = ["test response"] # Setup mock client mock_client = MagicMock() @@ -76,8 +76,8 @@ class TestVertexAIHTTPStatus201(unittest.TestCase): # Create a mock response with status code 201 mock_response = MagicMock() mock_response.status_code = 201 - mock_response.iter_lines = MagicMock() - mock_response.iter_lines.return_value = ["test response"] + mock_response.iter_text = MagicMock() + mock_response.iter_text.return_value = ["test response"] mock_post.return_value = mock_response # Call the make_sync_call function diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 009419e7a24..7ab2bc0e5d3 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -4,10 +4,11 @@ Mock tests for A2A endpoints. Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request. """ +import asyncio import json import socket import sys -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import AbstractContextManager, ExitStack from dataclasses import dataclass from typing import Final @@ -27,6 +28,21 @@ from litellm.types.caching import LiteLLMCacheType AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] +class _SSELineStream(httpx.AsyncByteStream): + def __init__(self, lines: tuple[str, ...], first_line_delay: float) -> None: + self._lines = lines + self._first_line_delay = first_line_delay + + async def __aiter__(self) -> AsyncIterator[bytes]: + await asyncio.sleep(self._first_line_delay) + for line in self._lines: + yield f"{line}\n".encode() + + +def _sse_upstream_response(*lines: str, first_line_delay: float = 0.0) -> httpx.Response: + return httpx.Response(200, stream=_SSELineStream(lines, first_line_delay)) + + @dataclass(frozen=True, slots=True) class CapturedAgentCall: request_id: object @@ -1016,13 +1032,7 @@ async def test_task_methods_extract_litellm_params_before_forwarding(method: str mock_http_response.json.return_value = upstream_response mock_http_response.is_success = True - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response('data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}') mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1057,8 +1067,8 @@ async def test_task_methods_extract_litellm_params_before_forwarding(method: str user_api_key_dict=user_api_key_dict, ) if method == "tasks/resubscribe": - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "task-1" in streamed if method == "tasks/resubscribe": forwarded_body = mock_async_client.build_request.call_args.kwargs["json"] @@ -1082,14 +1092,7 @@ async def test_subscribe_to_task_returns_sse_stream(): 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}', ] - async def fake_aiter_lines(): - for line in sse_lines: - yield line - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response(*sse_lines) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1139,13 +1142,9 @@ async def test_subscribe_to_task_calls_pre_call_hook(): mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"}) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response( + 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1","status":{"state":"completed"}}}' + ) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1192,9 +1191,9 @@ async def test_subscribe_to_task_calls_pre_call_hook(): ) assert response.media_type == "text/event-stream" - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "completed" in streamed mock_proxy_logging.pre_call_hook.assert_awaited_once() call_kwargs = mock_proxy_logging.pre_call_hook.await_args.kwargs assert call_kwargs.get("call_type") == "asend_message" @@ -1224,16 +1223,10 @@ async def test_subscribe_to_task_runs_post_call_streaming_guardrail(): mock_request = _make_request_mock("tasks/resubscribe", {"id": "task-1"}) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1", team_id="t1") - async def fake_aiter_lines(): - yield ( - 'data: {"jsonrpc":"2.0","id":"req-1","result":' - '{"kind":"message","parts":[{"kind":"text","text":"resubscribe-secret"}]}}' - ) - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response( + 'data: {"jsonrpc":"2.0","id":"req-1","result":' + '{"kind":"message","parts":[{"kind":"text","text":"resubscribe-secret"}]}}' + ) mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -1675,14 +1668,7 @@ async def test_pascal_method_names_normalize_to_wire_format(pascal_method: str, mock_http_response.is_success = True mock_http_response.raise_for_status = MagicMock() - async def _empty_aiter_lines(): - return - yield # make it an async generator - - mock_sse_resp = AsyncMock() - mock_sse_resp.is_success = True - mock_sse_resp.aiter_lines = _empty_aiter_lines - mock_sse_resp.aclose = AsyncMock() + mock_sse_resp = _sse_upstream_response() mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -2424,13 +2410,7 @@ async def test_task_methods_forward_trace_header(method: str): mock_http_response.json.return_value = upstream_response mock_http_response.is_success = True - async def fake_aiter_lines(): - yield 'data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}' - - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = fake_aiter_lines - mock_resp.aclose = AsyncMock() + mock_resp = _sse_upstream_response('data: {"jsonrpc":"2.0","id":"req-1","result":{"taskId":"task-1"}}') mock_async_client = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) @@ -2465,8 +2445,8 @@ async def test_task_methods_forward_trace_header(method: str): user_api_key_dict=user_api_key_dict, ) if method == "tasks/resubscribe": - async for _ in response.body_iterator: - pass + streamed = "".join([chunk async for chunk in response.body_iterator]) + assert "task-1" in streamed if method == "tasks/resubscribe": forwarded_headers = mock_async_client.build_request.call_args.kwargs["headers"] @@ -2721,17 +2701,12 @@ def test_served_version_falls_back_to_header_when_unconfigured(): assert _served_version(_agent(None), _request_with_a2a_header(None)) == "0.3" -def _sse_agent_handler(lines): - mock_resp = AsyncMock() - mock_resp.is_success = True - mock_resp.aiter_lines = lines - mock_resp.aclose = AsyncMock() - - mock_async_client = MagicMock() +def _sse_agent_handler(upstream_response: httpx.Response) -> MagicMock: + mock_async_client: Final = MagicMock() mock_async_client.build_request = MagicMock(return_value=MagicMock()) - mock_async_client.send = AsyncMock(return_value=mock_resp) + mock_async_client.send = AsyncMock(return_value=upstream_response) - mock_handler = MagicMock() + mock_handler: Final = MagicMock() mock_handler.client = mock_async_client return mock_handler @@ -2760,13 +2735,13 @@ async def test_forward_jsonrpc_sse_pings_while_the_upstream_agent_is_still_silen monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05) - async def _slow_lines(): - await asyncio.sleep(0.3) - yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}' + slow_upstream = _sse_upstream_response( + 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}', first_line_delay=0.3 + ) with patch( "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", - return_value=_sse_agent_handler(_slow_lines), + return_value=_sse_agent_handler(slow_upstream), ): response = await _resubscribe_response() assert response.headers["x-accel-buffering"] == "no" @@ -2788,12 +2763,11 @@ async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigure monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None) - async def _lines(): - yield 'data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}' - with patch( "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", - return_value=_sse_agent_handler(_lines), + return_value=_sse_agent_handler( + _sse_upstream_response('data: {"jsonrpc": "2.0", "id": "req-1", "result": {"kind": "task"}}') + ), ): response = await _resubscribe_response() assert "x-accel-buffering" not in response.headers @@ -2803,6 +2777,28 @@ async def test_forward_jsonrpc_sse_is_untouched_while_keepalives_are_unconfigure assert json.loads(chunks[-1].removeprefix("data: "))["result"]["kind"] == "task" +@pytest.mark.asyncio +async def test_forward_jsonrpc_sse_keeps_unicode_line_separators_inside_an_event(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", None) + separated_text: Final = "before\u2028between\u2029after" + event: Final = { + "jsonrpc": "2.0", + "id": "req-1", + "result": {"kind": "message", "parts": [{"kind": "text", "text": separated_text}]}, + } + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=_sse_agent_handler(_sse_upstream_response(f"data: {json.dumps(event, ensure_ascii=False)}")), + ): + response: Final = await _resubscribe_response() + chunks: Final = [chunk async for chunk in response.body_iterator] + + assert [json.loads(chunk.removeprefix("data: "))["result"]["parts"][0]["text"] for chunk in chunks] == [ + separated_text + ] + + async def _stream_message_response(): from litellm.proxy.agent_endpoints.a2a_endpoints import _handle_stream_message diff --git a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py index 85bb8304b95..5911fcdc073 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py +++ b/tests/unit/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -4,6 +4,7 @@ Tests for restamping the public model on Anthropic Messages streaming chunks. import json from collections.abc import AsyncIterator, Sequence +from typing import Final import pytest @@ -46,6 +47,28 @@ def test_restamps_sse_bytes_frame(): assert b"event: message_start" in restamped +def test_restamps_sse_bytes_frame_holding_a_line_separator() -> None: + payload: Final = { + "type": "message_start", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5-20251001", + "content": [{"type": "text", "text": "alpha\u2028beta"}], + }, + } + frame: Final = f"event: message_start\ndata: {json.dumps(payload, ensure_ascii=False)}\n\n".encode() + + restamped: Final = restamp_anthropic_stream_chunk_model(frame, "claude-auto-1") + + assert isinstance(restamped, bytes) + data_line: Final = next(line for line in restamped.decode().split("\n") if line.startswith("data:")) + message: Final = json.loads(data_line[len("data:") :])["message"] + assert message["model"] == "claude-auto-1" + assert message["content"] == [{"type": "text", "text": "alpha\u2028beta"}] + + def test_restamps_event_dict(): chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}} diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 092805a9cac..117428ee27b 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -5,7 +5,7 @@ import litellm.proxy import litellm.proxy.proxy_server -from typing import Dict, Final, List, Optional +from typing import Dict, Final, List, Mapping, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1367,7 +1367,9 @@ async def test_check_user_model_budget(user_id, user_model_max_budget, expected_ calls = [] class _Limiter: - async def is_user_within_model_budget(self, user_id, user_model_max_budget, model): + async def is_user_within_model_budget( + self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None + ) -> bool: calls.append((user_id, user_model_max_budget, model)) return True diff --git a/tests/unit/proxy/client/cli/test_auth_commands.py b/tests/unit/proxy/client/cli/test_auth_commands.py index cb1fcba9df7..5d5ebcf5dfd 100644 --- a/tests/unit/proxy/client/cli/test_auth_commands.py +++ b/tests/unit/proxy/client/cli/test_auth_commands.py @@ -3,13 +3,15 @@ import os import stat import time from pathlib import Path +from typing import Final from unittest.mock import Mock, patch import pytest +import requests import responses from click.testing import CliRunner +from responses import matchers -from litellm.constants import CLI_JWT_EXPIRATION_HOURS from litellm.litellm_core_utils.cli_keyring import ( DISABLE_KEYRING_ENV_VAR, KeyringDisabled, @@ -56,6 +58,20 @@ def _write_token_file(home: Path, *, key: str | None) -> None: _write_home_json(home, "token.json", payload if key is None else {**payload, "key": key}) +STALE_TOKEN_BASE_URL: Final = "https://stale-token.example.com" +STALE_TOKEN_TIMESTAMP: Final = 1_700_000_000.0 + + +def _register_models_answer(key: str, answer: int | Exception) -> None: + models_url: Final = f"{STALE_TOKEN_BASE_URL}/v1/models" + authorized: Final = [matchers.header_matcher({"Authorization": f"Bearer {key}"})] + if isinstance(answer, Exception): + responses.get(models_url, body=answer, match=authorized) + else: + responses.get(models_url, json={"data": []}, status=answer, match=authorized) + responses.get(models_url, status=401) + + def _secret_blob(base_url: str, key: str) -> str: return json.dumps({"base_url": base_url, "key": key, "jwt_token": ""}) @@ -696,22 +712,32 @@ class TestWhoamiCommand: assert "Not authenticated" in result.output assert "Run 'lite login'" in result.output - def test_whoami_old_token(self): - """Test whoami with old token showing warning""" - token_data = CliTokenRecord( + @pytest.mark.parametrize( + ("proxy_status", "expected_last_line_start"), + [ + (200, "Token age: "), + (401, "Token expired. Run 'lite login' again."), + ], + ) + @responses.activate + def test_whoami_lets_the_proxy_judge_a_token_older_than_the_client_window( + self, proxy_status: int, expected_last_line_start: str + ): + _register_models_answer("sk-live", proxy_status) + token_data: Final = CliTokenRecord( + base_url=STALE_TOKEN_BASE_URL, user_email="test@example.com", - user_id="test-user-123", - user_role="admin", key="sk-live", - timestamp=time.time() - (25 * 3600), + timestamp=STALE_TOKEN_TIMESTAMP, ) with patch("litellm.proxy.client.cli.commands.auth.load_cli_token", return_value=token_data): - result = self.runner.invoke(whoami) + result: Final = self.runner.invoke(whoami) - assert result.exit_code == 0 - assert "Authenticated" in result.output - assert "Warning: Token is more than 24 hours old" in result.output + assert result.exit_code == 0 + assert result.output.splitlines()[-1].startswith(expected_last_line_start) + assert "Warning" not in result.output + assert len(responses.calls) == 1 def test_whoami_missing_fields(self): """Test whoami with token missing some fields""" @@ -1024,51 +1050,93 @@ class TestPrintTokenCommand: assert result.exit_code == 0 assert result.output.strip() == "sk-matches" + @responses.activate def test_fresh_cached_key_printed_without_network_call(self): """A recently-issued key should be printed straight from cache -- no refresh call on every single invocation (apiKeyHelper gets called frequently).""" - with ( - patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", - return_value=CliTokenRecord( - base_url="http://localhost:4000", - key="sk-cached-fresh", - timestamp=time.time(), - ), + with patch( + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url="http://localhost:4000", + key="sk-cached-fresh", + timestamp=time.time(), ), - patch("requests.post") as mock_post, ): - result = self.runner.invoke(print_token, obj={}) + result: Final = self.runner.invoke(print_token, obj={}) assert result.exit_code == 0 assert result.output.strip() == "sk-cached-fresh" - mock_post.assert_not_called() + assert len(responses.calls) == 0 - def test_stale_key_fails_fast_without_network_call(self): - """There is no silent refresh: an expired cached key must fail - loudly (stderr, nonzero exit) telling the user to `lite login` - again, rather than making a network call or printing a dead key - that will just 401 Claude Code.""" - old_timestamp = time.time() - (CLI_JWT_EXPIRATION_HOURS + 1) * 3600 + @pytest.mark.parametrize( + ("proxy_answer", "expected_exit_code", "expected_stdout", "expected_stderr"), + [ + (200, 0, "sk-stale-key\n", ""), + (401, 1, "", "Token expired. Run 'lite login' again.\n"), + (403, 1, "", "Token expired. Run 'lite login' again.\n"), + (404, 0, "sk-stale-key\n", ""), + (500, 0, "sk-stale-key\n", ""), + (requests.ConnectionError("proxy unreachable"), 0, "sk-stale-key\n", ""), + ], + ) + @responses.activate + def test_key_older_than_the_client_window_is_judged_by_the_proxy( + self, + proxy_answer: int | Exception, + expected_exit_code: int, + expected_stdout: str, + expected_stderr: str, + ): + _register_models_answer("sk-stale-key", proxy_answer) - with ( - patch( - "litellm.proxy.client.cli.commands.auth.load_cli_token", - return_value=CliTokenRecord( - base_url="http://localhost:4000", - key="sk-stale-key", - timestamp=old_timestamp, - ), + with patch( + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url=STALE_TOKEN_BASE_URL, + key="sk-stale-key", + timestamp=STALE_TOKEN_TIMESTAMP, ), - patch("requests.post") as mock_post, ): - result = self.runner.invoke(print_token, obj={}) + result: Final = self.runner.invoke(print_token, obj={}) - assert result.exit_code != 0 - assert "sk-stale-key" not in result.output - assert "lite login" in result.output - mock_post.assert_not_called() + assert result.exit_code == expected_exit_code + assert result.stdout == expected_stdout + assert result.stderr == expected_stderr + assert len(responses.calls) == 1 + + @responses.activate + def test_proxy_check_waits_only_a_few_seconds_for_a_silent_proxy(self): + _register_models_answer("sk-stale-key", requests.ReadTimeout("proxy sent nothing")) + + with patch( + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord( + base_url=STALE_TOKEN_BASE_URL, + key="sk-stale-key", + timestamp=STALE_TOKEN_TIMESTAMP, + ), + ): + result: Final = self.runner.invoke(print_token, obj={}) + + timeout: Final = responses.calls[0].request.req_kwargs["timeout"] + assert result.exit_code == 0 + assert result.stdout == "sk-stale-key\n" + assert isinstance(timeout, int | float) + assert 0 < timeout <= 5 + + @responses.activate + def test_key_without_a_recorded_proxy_still_expires_after_the_client_window(self): + with patch( + "litellm.proxy.client.cli.commands.auth.load_cli_token", + return_value=CliTokenRecord(key="sk-old-format", timestamp=STALE_TOKEN_TIMESTAMP), + ): + result: Final = self.runner.invoke(print_token, obj={}) + + assert result.exit_code == 1 + assert result.stdout == "" + assert result.stderr == "Token expired. Run 'lite login' again.\n" + assert len(responses.calls) == 0 class TestPrintTokenWithConfigFile: @@ -2019,18 +2087,27 @@ class TestPkcePrintToken: assert sum(len(session.posts) for session in _FakeSession.instances) == 1 assert save.call_count == 1 - def test_print_token_for_an_expired_classic_token_makes_no_request(self): + @responses.activate + def test_print_token_trusts_a_recorded_expiry_without_asking_the_proxy(self): + _register_models_answer("sk-classic", 200) with ( patch( "litellm.proxy.client.cli.commands.auth.load_token", - return_value={"key": "sk-classic", "timestamp": time.time() - (CLI_JWT_EXPIRATION_HOURS + 1) * 3600}, + return_value={ + "key": "sk-classic", + "base_url": STALE_TOKEN_BASE_URL, + "timestamp": STALE_TOKEN_TIMESTAMP, + "expires_at": STALE_TOKEN_TIMESTAMP + 3600, + }, ), patch("litellm.proxy.client.cli.commands.auth.requests.Session", _FakeSession), ): - result = self.runner.invoke(print_token, obj={}) + result: Final = self.runner.invoke(print_token, obj={}) assert result.exit_code == 1 - assert "Token expired" in result.output + assert result.stdout == "" + assert result.stderr == "Token expired. Run 'lite login' again.\n" + assert len(responses.calls) == 0 assert _FakeSession.instances == [] diff --git a/tests/unit/proxy/db/test_accounting_ledger.py b/tests/unit/proxy/db/test_accounting_ledger.py new file mode 100644 index 00000000000..3916fe869e6 --- /dev/null +++ b/tests/unit/proxy/db/test_accounting_ledger.py @@ -0,0 +1,103 @@ +import json +from datetime import datetime, timezone +from decimal import Decimal +from typing import Final, Literal, NoReturn + +import pytest + +from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome +from litellm.proxy.db.accounting_ledger import AccountingLedger +from litellm.types.utils import CallTypes + + +def _no_transaction() -> NoReturn: + raise AssertionError("Invalid accounting input must not open a database transaction") + + +@pytest.fixture +def event() -> InboundUsageEvent: + return InboundUsageEvent( + event_id="caller", + request_id="caller", + call_type=CallTypes.completion, + started_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + outcome=Outcome.SUCCEEDED, + dimensions=(), + metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("envelope", ["{}", "[]", '{"event":null}', '{"event":{}}', "not json"]) +async def test_invalid_envelopes_fail_before_database_access(event: InboundUsageEvent, envelope: str) -> None: + with pytest.raises(ValueError, match=r"Accounting envelope|validation error"): + await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [0, -1, True, 1.5]) +@pytest.mark.parametrize("field", ["version", "schema_version"]) +async def test_invalid_versions_fail_before_database_access( + event: InboundUsageEvent, version: int | float, field: Literal["version", "schema_version"] +) -> None: + with pytest.raises(ValueError, match="positive integers"): + await AccountingLedger(_no_transaction).insert_event( + event, + (), + envelope="{}", + **{field: version}, # pyright: ignore[reportArgumentType] # Exercise invalid runtime versions + ) + + +@pytest.mark.asyncio +async def test_constructed_invalid_events_are_revalidated(event: InboundUsageEvent) -> None: + invalid: Final = event.model_copy(update={"event_id": ""}) + with pytest.raises(ValueError, match="event_id"): + await AccountingLedger(_no_transaction).insert_event(invalid, (), envelope="{}") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit", [0, -1, 1001, True, 1.5]) +async def test_retention_rejects_invalid_batch_bounds(limit: int | float) -> None: + with pytest.raises(ValueError, match="batch size"): + await AccountingLedger(_no_transaction).purge_completed( + limit=limit # pyright: ignore[reportArgumentType] # Exercise invalid runtime batch sizes + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("missing", ["both", "redelivery", "recovery"]) +async def test_retention_requires_both_safety_watermarks(missing: Literal["both", "redelivery", "recovery"]) -> None: + cutoff: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + deleted: Final = await AccountingLedger(_no_transaction).purge_completed( + redelivery_before=cutoff if missing == "recovery" else None, + recovery_before=cutoff if missing == "redelivery" else None, + ) + assert deleted == 0 + + +@pytest.mark.asyncio +async def test_retention_rejects_ambiguous_naive_watermarks() -> None: + with pytest.raises(ValueError, match="timezone"): + await AccountingLedger(_no_transaction).purge_completed( + redelivery_before=datetime(2026, 1, 1), recovery_before=datetime(2026, 1, 1, tzinfo=timezone.utc) + ) + + +@pytest.mark.asyncio +async def test_nonfinite_context_is_rejected_before_database_access(event: InboundUsageEvent) -> None: + envelope: Final = json.dumps({"event": event.model_dump(mode="json"), "context": {"price": float("nan")}}) + with pytest.raises(ValueError, match="Out of range float"): + await AccountingLedger(_no_transaction).insert_event(event, (), envelope=envelope) + + +@pytest.mark.asyncio +async def test_unknown_delivery_target_is_rejected_before_database_access(event: InboundUsageEvent) -> None: + envelope: Final = json.dumps({"event": event.model_dump(mode="json")}) + with pytest.raises(ValueError, match="not a valid DeliveryTarget"): + await AccountingLedger(_no_transaction).insert_event( + event, + ("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target + envelope=envelope, + ) diff --git a/tests/unit/proxy/db/test_budget_window_spend_writer.py b/tests/unit/proxy/db/test_budget_window_spend_writer.py index 6fc438feee8..27f1472c183 100644 --- a/tests/unit/proxy/db/test_budget_window_spend_writer.py +++ b/tests/unit/proxy/db/test_budget_window_spend_writer.py @@ -9,6 +9,7 @@ import pytest from litellm.proxy.db.budget_window_spend_writer import ( WindowSeedTotals, commit_window_spend_updates, + open_window_spend_row, roll_window_spend_row, spend_logs_seed_totals, ) @@ -405,6 +406,29 @@ async def test_roll_window_spend_row_is_conditional_on_the_stored_window_being_o assert params[:4] == ("team", "t1", "30d", datetime(2026, 8, 31)) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("replaces_window_start", "replaced_param"), [(None, None), (WINDOW_A, datetime(2026, 8, 1))] +) +async def test_open_window_spend_row_moves_only_the_window_the_caller_saw_end(replaces_window_start, replaced_param): + db = _FakeDB() + + await open_window_spend_row( + prisma_client=_FakePrismaClient(db), + entity_type="model_max_budget", + entity_id="k1", + window_duration="1d", + window_start=WINDOW_B, + replaces_window_start=replaces_window_start, + ) + + ((query, params),) = db.execute_raw_calls + normalized = " ".join(query.split()) + assert "ON CONFLICT (entity_type, entity_id, window_duration) DO UPDATE" in normalized + assert normalized.endswith("WHERE \"LiteLLM_BudgetWindowSpend\".window_start <= ($5::timestamptz AT TIME ZONE 'UTC')") + assert params == ("model_max_budget", "k1", "1d", datetime(2026, 8, 31), replaced_param) + + @pytest.mark.asyncio async def test_seed_receives_the_batch_earliest_start_as_its_cutoff(): db = _FakeDB(existing_rows=[]) diff --git a/tests/unit/proxy/decisions_endpoints/test_endpoints.py b/tests/unit/proxy/decisions_endpoints/test_endpoints.py index a0a79f9b7cf..c67d1012475 100644 --- a/tests/unit/proxy/decisions_endpoints/test_endpoints.py +++ b/tests/unit/proxy/decisions_endpoints/test_endpoints.py @@ -1834,3 +1834,56 @@ def test_proxy_routes_lm_studio_decisions_with_optional_authentication( else [{"type": "predicate", "name": "is_defect", "probability": 0.9}] ) assert response.json()["answers"] == expected + + +@pytest.mark.parametrize("provider", ("ollama", "ollama_chat")) +@pytest.mark.parametrize("endpoint", ("/v1/systemone", "/v1/decisions")) +@pytest.mark.parametrize("api_key", (None, "deployment-key")) +def test_proxy_routes_ollama_decisions_with_optional_authentication( + client: TestClient, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + provider: str, + endpoint: str, + api_key: str | None, +) -> None: + monkeypatch.delenv("OLLAMA_API_KEY", raising=False) + monkeypatch.delenv("OLLAMA_API_BASE", raising=False) + deployment_router: Final = litellm.Router( + model_list=[ + { + "model_name": "local-decision", + "litellm_params": { + "model": f"{provider}/d3:4b", + "api_base": "http://ollama.local:11434/v1", + "api_key": api_key, + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", deployment_router) + upstream: Final = respx_mock.post("http://ollama.local:11434/v1/systemone").respond(json=_RESPONSE) + fields: Final = ( + {"state": "review", "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}} + if endpoint.endswith("systemone") + else { + "input": "review", + "questions": [{"type": "predicate", "name": "is_defect", "instructions": "Is this a defect?"}], + } + ) + + response: Final = client.post(endpoint, json={"model": "local-decision", **fields}) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == { + "model": "d3:4b", + "state": "review", + "questions": {"is_defect": {"type": "noul", "instructions": "Is this a defect?"}}, + } + assert upstream.calls.last.request.headers.get("authorization") == (f"Bearer {api_key}" if api_key else None) + expected: Final = ( + _RESPONSE["answers"] + if endpoint.endswith("systemone") + else [{"type": "predicate", "name": "is_defect", "probability": 0.9}] + ) + assert response.json()["answers"] == expected diff --git a/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py b/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py index 79062f15eca..819e63dfd59 100644 --- a/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_gemini_agents_endpoints.py @@ -8,12 +8,18 @@ query parameter. Flat credential query params (e.g. ?api_key=…) are no longer accepted — they would appear in server logs. """ +import asyncio import json +from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx from fastapi import Request from fastapi.datastructures import Headers, QueryParams +from fastapi.testclient import TestClient from litellm.proxy.google_endpoints.agents_endpoints import ( @@ -470,3 +476,41 @@ def test_validate_environment_env_fallback_when_no_api_base_override(monkeypatch cfg = GeminiAgentsConfig() headers = cfg.validate_environment(headers={}, litellm_params={}) assert headers["x-goog-api-key"] == "AIzaFromEnv" + + +_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "test_configs" / "test_bad_config.yaml" + + +@pytest.fixture +def proxy_client(monkeypatch: pytest.MonkeyPatch, httpx_transport: None) -> TestClient: + from litellm.proxy.proxy_server import app, cleanup_router_config_variables, initialize + + monkeypatch.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") + cleanup_router_config_variables() + asyncio.run(initialize(config=str(_PROXY_CONFIG))) + return TestClient(app) + + +def test_create_agent_route_forwards_googles_create_body(proxy_client: TestClient, respx_mock: respx.MockRouter): + google_body: Final = { + "id": "sentinel", + "base_agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + "system_instruction": "Always answer in exactly three words.", + "base_environment": {"type": "remote"}, + "description": "Watches the queue.", + "tools": [{"type": "code_execution"}], + } + upstream: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/agents").mock( + return_value=httpx.Response(200, json=google_body) + ) + + response: Final = proxy_client.post( + "/v1beta/agents", + json={**google_body, "litellm_params_template": {"api_key": "AIza-caller"}}, + ) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == google_body + assert upstream.calls.last.request.headers["x-goog-api-key"] == "AIza-caller" + assert response.json()["id"] == "sentinel" diff --git a/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py b/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py index f3cec320532..2a2a00d322e 100644 --- a/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py +++ b/tests/unit/proxy/google_endpoints/test_interactions_agent_param.py @@ -6,9 +6,16 @@ The proxy must not pass the agent name as ``model`` or LiteLLM may route to openai/* wildcards instead of Gemini interactions. """ +import asyncio +import json +from pathlib import Path +from typing import Final from unittest.mock import MagicMock, patch +import httpx import pytest +import respx +from fastapi.testclient import TestClient class TestInteractionsAgentParameter: @@ -101,3 +108,37 @@ class TestInteractionsAgentOnlyProviderRouting: mock_get_llm_provider.assert_not_called() assert mock_handler.create_interaction.call_args.kwargs["model"] is None + + +_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "test_configs" / "test_bad_config.yaml" + + +@pytest.fixture +def proxy_client(monkeypatch: pytest.MonkeyPatch, httpx_transport: None) -> TestClient: + from litellm.proxy.proxy_server import app, cleanup_router_config_variables, initialize + + monkeypatch.setenv("LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", "true") + cleanup_router_config_variables() + asyncio.run(initialize(config=str(_PROXY_CONFIG))) + return TestClient(app) + + +@pytest.mark.usefixtures("fake_provider_credentials") +def test_create_interaction_route_forwards_agent_config(proxy_client: TestClient, respx_mock: respx.MockRouter): + request_body: Final = { + "agent": "antigravity-preview-09-2026", + "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, + "input": "Say hi in one word.", + "environment": "remote", + } + upstream: Final = respx_mock.post("https://generativelanguage.googleapis.com/v1beta/interactions").mock( + return_value=httpx.Response( + 200, + json={"id": "interaction-agent", "object": "interaction", "status": "completed"}, + ) + ) + + response: Final = proxy_client.post("/v1beta/interactions", json=request_body) + + assert response.status_code == 200, response.text + assert json.loads(upstream.calls.last.request.content) == request_body diff --git a/tests/unit/proxy/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py index a668d5f609d..21d96a3ea97 100644 --- a/tests/unit/proxy/guardrails/test_content_utils.py +++ b/tests/unit/proxy/guardrails/test_content_utils.py @@ -2,6 +2,7 @@ import json from functools import reduce +from typing import Final import pytest from pydantic import JsonValue @@ -924,6 +925,34 @@ def test_native_writer_visits_each_tool_and_system_slot(data): assert "payload" not in json.dumps(data) +def test_skip_tool_message_leaves_tool_results_and_masks_the_rest() -> None: + data: Final[dict[str, object]] = { + "messages": [ + {"role": "user", "content": "secret"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tool", "name": "read", "input": {"env": "secret"}}], + "tool_calls": [{"id": "call", "function": {"name": "read", "arguments": '{"env": "secret"}'}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tool", "content": "secret"}]}, + {"role": "tool", "tool_call_id": "call", "content": "secret"}, + ] + } + assert walk_user_text(data, lambda text: text.replace("secret", "[MASKED]"), skip_tool_message=True) == 3 + assert data == { + "messages": [ + {"role": "user", "content": "[MASKED]"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "tool", "name": "read", "input": {"env": "[MASKED]"}}], + "tool_calls": [{"id": "call", "function": {"name": "read", "arguments": '{"env": "[MASKED]"}'}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tool", "content": "secret"}]}, + {"role": "tool", "tool_call_id": "call", "content": "secret"}, + ] + } + + @pytest.mark.parametrize( "part", ( diff --git a/tests/unit/proxy/hooks/test_model_budget_shared_windows.py b/tests/unit/proxy/hooks/test_model_budget_shared_windows.py new file mode 100644 index 00000000000..7002f602fca --- /dev/null +++ b/tests/unit/proxy/hooks/test_model_budget_shared_windows.py @@ -0,0 +1,199 @@ +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from typing import Final + +import pytest + +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction +from litellm.proxy.hooks.model_budget_shared_windows import DatabaseSharedWindows, SharedWindow + +SPEND_KEY: Final = "virtual_key_spend:key-hash:claude-opus*:1d" +WINDOW_START: Final = datetime(2026, 10, 20, 9, 30, tzinfo=timezone.utc) + + +class _Table: + def __init__(self, row: object | None = None, error: Exception | None = None) -> None: + self.row = row + self.error = error + self.reads = 0 + + async def find_unique(self, where: object) -> object | None: + self.reads = self.reads + 1 + if self.error is not None: + raise self.error + return self.row + + +class _Database: + def __init__(self, table: _Table) -> None: + self.litellm_budgetwindowspend = table + self.opened_with: tuple[object, ...] = () + + async def execute_raw(self, query: str, *params: object) -> int: + self.opened_with = params + self.litellm_budgetwindowspend.row = SimpleNamespace(window_start=params[3], spend=0.0) + return 1 + + +class _OutracedDatabase: + def __init__(self, table: _Table, winner: object) -> None: + self.litellm_budgetwindowspend = table + self._winner = winner + + async def execute_raw(self, query: str, *params: object) -> int: + self.litellm_budgetwindowspend.row = self._winner + return 0 + + +class _Sink: + def __init__(self) -> None: + self.transactions: tuple[WindowSpendTransaction, ...] = () + + async def __call__(self, transaction: WindowSpendTransaction) -> None: + self.transactions = (*self.transactions, transaction) + + +class _LaggingReplica: + def __init__(self, writer: _Database) -> None: + self.litellm_budgetwindowspend = _Table() + self._writer = writer + + async def execute_raw(self, query: str, *params: object) -> int: + return await self._writer.execute_raw(query, *params) + + +def _database_windows(table: _Table | None, sink: _Sink, lagging_replica: bool = False) -> DatabaseSharedWindows: + writer: Final = None if table is None else _Database(table) + prisma: Final = ( + None + if writer is None + else SimpleNamespace(db=_LaggingReplica(writer) if lagging_replica else writer, writer_db=writer) + ) + return DatabaseSharedWindows( + prisma_client_source=lambda: prisma, # pyright: ignore[reportArgumentType] # duck-typed prisma double + sink=sink, + ) + + +def test_a_window_is_open_until_its_period_has_passed(): + window: Final = SharedWindow(started_at=WINDOW_START.timestamp(), spend=1.0) + day: Final = 24 * 60 * 60 + + assert window.is_open(budget_duration="1d", now=WINDOW_START.timestamp() + day - 1) + assert not window.is_open(budget_duration="1d", now=WINDOW_START.timestamp() + day) + + +@pytest.mark.asyncio +async def test_read_returns_the_stored_window_and_rereads_only_after_the_cache_expires(): + table: Final = _Table(row=SimpleNamespace(window_start=WINDOW_START.replace(tzinfo=None), spend=2.5)) + windows: Final = _database_windows(table, _Sink()) + + first: Final = await windows.read(spend_key=SPEND_KEY, budget_duration="1d") + second: Final = await windows.read(spend_key=SPEND_KEY, budget_duration="1d") + + assert first == SharedWindow(started_at=WINDOW_START.timestamp(), spend=2.5) + assert second == first + assert table.reads == 1 + + +@pytest.mark.asyncio +async def test_an_unreadable_table_reads_as_no_shared_window(): + table: Final = _Table(error=RuntimeError("database is down")) + windows: Final = _database_windows(table, _Sink()) + + assert await windows.read(spend_key=SPEND_KEY, budget_duration="1d") is None + + +@pytest.mark.asyncio +async def test_add_queues_the_charge_for_the_window_it_belongs_to(): + sink: Final = _Sink() + windows: Final = _database_windows(_Table(), sink) + + await windows.add(spend_key=SPEND_KEY, budget_duration="1d", started_at=WINDOW_START.timestamp(), spend=0.25) + + assert len(sink.transactions) == 1 + queued: Final = sink.transactions[0] + assert (queued["entity_id"], queued["window_duration"], queued["spend"]) == (SPEND_KEY, "1d", 0.25) + assert datetime.fromisoformat(queued["window_start"]).replace(tzinfo=timezone.utc) == WINDOW_START + + +@pytest.mark.asyncio +async def test_nothing_is_queued_without_a_database(): + sink: Final = _Sink() + windows: Final = _database_windows(None, sink) + + await windows.add(spend_key=SPEND_KEY, budget_duration="1d", started_at=WINDOW_START.timestamp(), spend=0.25) + + assert sink.transactions == () + assert await windows.read(spend_key=SPEND_KEY, budget_duration="1d") is None + + +@pytest.mark.asyncio +async def test_open_joins_the_window_another_replica_already_started(): + table: Final = _Table(row=SimpleNamespace(window_start=WINDOW_START.replace(tzinfo=None), spend=2.5)) + windows: Final = _database_windows(table, _Sink()) + + started_at: Final = await windows.open( + spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp() + 3600 + ) + + assert started_at == WINDOW_START.timestamp() + assert table.row.spend == 2.5 + + +@pytest.mark.asyncio +async def test_open_reads_the_window_from_the_writer_while_the_read_replica_lags(): + table: Final = _Table(row=SimpleNamespace(window_start=WINDOW_START.replace(tzinfo=None), spend=2.5)) + windows: Final = _database_windows(table, _Sink(), lagging_replica=True) + + started_at: Final = await windows.open( + spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp() + 3600 + ) + + assert started_at == WINDOW_START.timestamp() + + +@pytest.mark.asyncio +async def test_open_joins_the_window_a_racing_replica_opened_first(): + winner: Final = SimpleNamespace(window_start=WINDOW_START.replace(tzinfo=None), spend=2.5) + database: Final = _OutracedDatabase(_Table(), winner) + windows: Final = DatabaseSharedWindows( + prisma_client_source=lambda: SimpleNamespace(db=database, writer_db=database), # pyright: ignore[reportArgumentType] # duck-typed prisma double + sink=_Sink(), + ) + + started_at: Final = await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp() + 60) + + assert started_at == WINDOW_START.timestamp() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored_start", [None, WINDOW_START - timedelta(days=2)]) +async def test_open_starts_a_window_at_now_when_none_is_open(stored_start: datetime | None): + table: Final = _Table( + row=None if stored_start is None else SimpleNamespace(window_start=stored_start.replace(tzinfo=None), spend=9.0) + ) + windows: Final = _database_windows(table, _Sink()) + + started_at: Final = await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp()) + + assert started_at == WINDOW_START.timestamp() + assert await windows.read(spend_key=SPEND_KEY, budget_duration="1d") == SharedWindow( + started_at=WINDOW_START.timestamp(), spend=0.0 + ) + + +@pytest.mark.asyncio +async def test_open_starts_this_replicas_own_window_without_a_database(): + windows: Final = _database_windows(None, _Sink()) + + assert await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp()) == ( + WINDOW_START.timestamp() + ) + + +@pytest.mark.asyncio +async def test_open_reports_no_shared_window_while_the_database_is_unreachable(): + windows: Final = _database_windows(_Table(error=RuntimeError("database is down")), _Sink()) + + assert await windows.open(spend_key=SPEND_KEY, budget_duration="1d", now=WINDOW_START.timestamp()) is None diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index 2678f4205d5..f5a2ccf7b76 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -1,19 +1,30 @@ import asyncio +import functools +import re import time from collections.abc import Callable, Mapping -from types import MappingProxyType +from dataclasses import dataclass +from datetime import datetime, timezone +from types import MappingProxyType, SimpleNamespace from typing import Final import pytest import litellm from litellm.caching.caching import DualCache -from litellm.proxy._types import Litellm_EntityType +from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth +from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction +from litellm.proxy.hooks.model_budget_shared_windows import DatabaseSharedWindows, SharedWindows from litellm.proxy.hooks.model_max_budget_limiter import ( PROXY_VirtualKeyModelMaxBudgetLimiter, + build_model_max_budget_usage, + model_access_groups_of, model_budget_spend_cache_key, + resolve_model_budget, + routed_models_of, team_member_budget_entity_id, ) +from litellm.router import Router from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import LiteLLMBatch, Usage @@ -74,10 +85,14 @@ def _event( } +async def _log_success( + limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, event: dict[str, object], response_obj: object = None +) -> None: + await limiter.async_log_success_event(event, response_obj=response_obj, start_time=None, end_time=None) + + async def _poll(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMBatch, response_cost: float) -> None: - await limiter.async_log_success_event( - _event("aretrieve_batch", response_cost), response_obj=batch, start_time=None, end_time=None - ) + await _log_success(limiter, _event("aretrieve_batch", response_cost), response_obj=batch) async def _chat( @@ -86,7 +101,8 @@ async def _chat( user_id: str = USER_ID, team_id: str = TEAM_ID, ) -> None: - await limiter.async_log_success_event( + await _log_success( + limiter, _event( "acompletion", CHAT_COST, @@ -94,9 +110,6 @@ async def _chat( user_id=user_id, team_id=team_id, ), - response_obj=None, - start_time=None, - end_time=None, ) @@ -311,3 +324,568 @@ async def test_a_batch_polled_within_every_budget_window_is_never_charged_again( await _poll(limiter, finished, BATCH_COST) assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + + +OPUS_GROUP: Final = "claude-opus" +OPUS_GROUP_MEMBERS: Final = MappingProxyType( + {"claude-opus-4-8": frozenset({OPUS_GROUP}), "claude-opus-5-5": frozenset({OPUS_GROUP})} +) +FAMILY_CAP: Final = {"budget_limit": CHAT_COST, "time_period": "1d"} # mutable-ok: BudgetConfig validates dicts only +ROOMY_CAP: Final = { + "budget_limit": CHAT_COST * 100, + "time_period": "1d", +} # mutable-ok: BudgetConfig validates dicts only +WINDOW_EPOCH: Final = datetime(2026, 10, 20, 9, 30, tzinfo=timezone.utc).timestamp() + + +def _model_access_groups( + groups_by_model: Mapping[str, frozenset[str]], +) -> Callable[[str, str | None], frozenset[str]]: + return lambda model, _team_id: groups_by_model.get(model, frozenset()) + + +@dataclass(frozen=True, slots=True) +class _WindowRow: + window_start: datetime + spend: float + + +class _WindowSpendDatabase: + def __init__(self) -> None: + self.rows: Mapping[str, _WindowRow] = MappingProxyType({}) + self.queued: tuple[WindowSpendTransaction, ...] = () + self.reachable = True + + @property + def litellm_budgetwindowspend(self) -> "_WindowSpendDatabase": + return self + + async def find_unique(self, where: Mapping[str, Mapping[str, str]]) -> _WindowRow | None: + await asyncio.sleep(0) + self._answer() + return self.rows.get(where["entity_type_entity_id_window_duration"]["entity_id"]) + + async def execute_raw( + self, + query: str, + entity_type: str, + entity_id: str, + window_duration: str, + window_start: datetime, + replaces_window_start: datetime | None, + ) -> int: + await asyncio.sleep(0) + self._answer() + stored: Final = self.rows.get(entity_id) + if stored is not None and (replaces_window_start is None or stored.window_start > replaces_window_start): + return 0 + self.rows = MappingProxyType({**self.rows, entity_id: _WindowRow(window_start=window_start, spend=0.0)}) + return 1 + + async def queue(self, transaction: WindowSpendTransaction) -> None: + self.queued = (*self.queued, transaction) + + def flush(self) -> None: + self.rows = functools.reduce(_upsert_window_spend, self.queued, self.rows) + self.queued = () + + def replica(self) -> DatabaseSharedWindows: + return DatabaseSharedWindows( + prisma_client_source=lambda: SimpleNamespace(db=self, writer_db=self), # pyright: ignore[reportArgumentType] # duck-typed prisma double + sink=self.queue, + ) + + def _answer(self) -> None: + if not self.reachable: + raise ConnectionError("database is unreachable") + + +def _upsert_window_spend( + rows: Mapping[str, _WindowRow], transaction: WindowSpendTransaction +) -> Mapping[str, _WindowRow]: + window_start: Final = datetime.fromisoformat(transaction["window_start"]) + stored: Final = rows.get(transaction["entity_id"]) + row: Final = ( + _WindowRow(window_start=stored.window_start, spend=stored.spend + transaction["spend"]) + if stored is not None and stored.window_start >= window_start + else _WindowRow(window_start=window_start, spend=transaction["spend"]) + ) + return MappingProxyType({**rows, transaction["entity_id"]: row}) + + +def _family_limiter( + groups_by_model: Mapping[str, frozenset[str]] = OPUS_GROUP_MEMBERS, + dual_cache: DualCache | None = None, + shared_windows: SharedWindows | None = None, + clock: Callable[[], float] = time.time, +) -> PROXY_VirtualKeyModelMaxBudgetLimiter: + return PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=dual_cache or DualCache(), + clock=clock, + model_access_groups=_model_access_groups(groups_by_model), + shared_windows=shared_windows or _WindowSpendDatabase().replica(), + ) + + +def _key(model_max_budget: Mapping[str, object]) -> UserAPIKeyAuth: + return UserAPIKeyAuth(token=KEY_HASH, key_alias="opus-family", model_max_budget=dict(model_max_budget)) + + +async def _spend_on( + limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, model: str, model_max_budget: Mapping[str, object] +) -> None: + await _log_success( + limiter, + { + "call_type": "acompletion", + "standard_logging_object": { + "call_type": "acompletion", + "response_cost": CHAT_COST, + "model": f"anthropic/{model}", + "model_group": model, + "metadata": {"user_api_key_hash": KEY_HASH}, + }, + "litellm_params": {"metadata": {"user_api_key_model_max_budget": dict(model_max_budget)}}, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("family_entry", ["claude-opus*", OPUS_GROUP]) +async def test_every_opus_version_draws_from_one_family_entry(family_entry: str): + budget: Final = {family_entry: FAMILY_CAP} + limiter: Final = _family_limiter() + assert await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5") + + await _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget) + + with pytest.raises( + litellm.BudgetExceededError, match=rf"model=claude-opus-5-5, budget_model={re.escape(family_entry)}$" + ): + await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5") + assert await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-sonnet-5") + + +@pytest.mark.asyncio +async def test_a_model_added_to_the_access_group_later_is_covered_with_no_key_change(): + budget: Final = {OPUS_GROUP: FAMILY_CAP} + shared_cache: Final = DualCache() + before_group_change: Final = _family_limiter(dual_cache=shared_cache) + after_group_change: Final = _family_limiter( + groups_by_model={**OPUS_GROUP_MEMBERS, "claude-opus-6": frozenset({OPUS_GROUP})}, dual_cache=shared_cache + ) + + await _spend_on(before_group_change, model="claude-opus-4-8", model_max_budget=budget) + + assert await before_group_change.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-6") + with pytest.raises(litellm.BudgetExceededError): + await after_group_change.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-6") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("budget", "checked_model", "blocked"), + [ + ({"claude-opus-4-8": ROOMY_CAP, "claude-opus*": FAMILY_CAP}, "claude-opus-4-8", False), + ({"claude-opus-4-8": ROOMY_CAP, "claude-opus*": FAMILY_CAP}, "claude-opus-5-5", False), + ({OPUS_GROUP: FAMILY_CAP, "claude-opus*": ROOMY_CAP}, "claude-opus-5-5", True), + ({"claude-opus-4*": FAMILY_CAP, "claude-opus*": ROOMY_CAP}, "claude-opus-4-7", True), + ({"claude-opus-4*": FAMILY_CAP, "claude-opus*": ROOMY_CAP}, "claude-opus-5-5", False), + ({"claude-opus-4*": {"budget_limit": CHAT_COST}, "claude-opus*": FAMILY_CAP}, "claude-opus-5-5", True), + ({"claude-opus-4*": FAMILY_CAP, "c" + "*" * 15: ROOMY_CAP}, "claude-opus-4-7", True), + ], +) +async def test_the_most_specific_entry_takes_the_spend_and_gates_the_request( + budget: Mapping[str, object], checked_model: str, blocked: bool +): + limiter: Final = _family_limiter() + + await _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget) + + if blocked: + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model=checked_model) + else: + assert await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model=checked_model) + + +@pytest.mark.asyncio +async def test_an_exact_entry_refusal_reads_as_before_for_a_provider_prefixed_model(): + budget: Final = {"claude-opus-4-8": FAMILY_CAP} + limiter: Final = _family_limiter() + + await _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget) + + with pytest.raises(litellm.BudgetExceededError, match=r"model=anthropic/claude-opus-4-8$"): + await limiter.is_key_within_model_budget(user_api_key_dict=_key(budget), model="anthropic/claude-opus-4-8") + + +async def _spend_on_team_key( + limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, + model: str, + key_model_max_budget: Mapping[str, object], + team_model_max_budget: Mapping[str, object], +) -> None: + await _log_success( + limiter, + { + "call_type": "acompletion", + "standard_logging_object": { + "call_type": "acompletion", + "response_cost": CHAT_COST, + "model": f"anthropic/{model}", + "model_group": model, + "metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_team_id": TEAM_ID}, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": dict(key_model_max_budget), + "user_api_key_team_model_max_budget": dict(team_model_max_budget), + } + }, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("team_model_max_budget", "team_blocks"), + [({"claude-opus*": FAMILY_CAP}, True), ({OPUS_GROUP: FAMILY_CAP}, True), ({"claude-opus-4-8": FAMILY_CAP}, False)], +) +async def test_a_key_entry_named_after_the_model_replaces_only_the_team_entry_of_that_name( + team_model_max_budget: Mapping[str, object], team_blocks: bool +): + key_model_max_budget: Final = {"claude-opus-4-8": ROOMY_CAP} + limiter: Final = _family_limiter() + + await _spend_on_team_key( + limiter, + model="claude-opus-4-8", + key_model_max_budget=key_model_max_budget, + team_model_max_budget=team_model_max_budget, + ) + team_check: Final = limiter.is_team_within_model_budget( + team_id=TEAM_ID, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="claude-opus-4-8", + ) + + if team_blocks: + with pytest.raises(litellm.BudgetExceededError, match=r"budget_model="): + await team_check + else: + assert await team_check + + +TEAM_OPUS_GROUP: Final = "team-opus" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget_entry", [TEAM_OPUS_GROUP, "claude-opus*"]) +async def test_a_team_model_is_gated_and_charged_under_its_public_name(budget_entry: str): + router: Final = Router( + model_list=[ + { + "model_name": f"model_name_{TEAM_ID}_opus", + "litellm_params": {"model": "anthropic/claude-opus-4-8", "api_key": "sk-test"}, + "model_info": { + "id": "team-opus-deployment", + "team_id": TEAM_ID, + "team_public_model_name": "claude-opus-4-8", + "access_groups": [TEAM_OPUS_GROUP], + }, + } + ] + ) + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=DualCache(), + model_access_groups=lambda model, team_id: model_access_groups_of(router=router, model=model, team_id=team_id), + shared_windows=_WindowSpendDatabase().replica(), + ) + budget: Final = {budget_entry: FAMILY_CAP} + team_key: Final = UserAPIKeyAuth(token=KEY_HASH, team_id=TEAM_ID, model_max_budget=budget) + assert await limiter.is_key_within_model_budget(user_api_key_dict=team_key, model="claude-opus-4-8") + + await _log_success( + limiter, + { + "call_type": "acompletion", + "standard_logging_object": { + "call_type": "acompletion", + "response_cost": CHAT_COST, + "model": "anthropic/claude-opus-4-8", + "model_group": f"model_name_{TEAM_ID}_opus", + "metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_team_id": TEAM_ID}, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": budget, + "model_info": router.get_deployment(model_id="team-opus-deployment").model_info.model_dump(), + } + }, + }, + ) + + with pytest.raises(litellm.BudgetExceededError, match=rf"budget_model={re.escape(budget_entry)}$"): + await limiter.is_key_within_model_budget(user_api_key_dict=team_key, model="claude-opus-4-8") + + +@pytest.mark.asyncio +async def test_spend_charged_while_the_database_is_unreachable_reaches_the_shared_window_without_wiping_it(): + budget: Final = {"claude-opus*": ROOMY_CAP} + database: Final = _WindowSpendDatabase() + clock: Final = _Clock() + clock.advance(WINDOW_EPOCH) + replica_a: Final = _family_limiter(shared_windows=database.replica(), clock=clock.now) + replica_b: Final = _family_limiter(shared_windows=database.replica(), clock=clock.now) + + await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget) + database.flush() + clock.advance(60) + database.reachable = False + await _spend_on(replica_b, model="claude-opus-4-8", model_max_budget=budget) + database.reachable = True + database.flush() + clock.advance(60) + await _spend_on(replica_b, model="claude-opus-5-5", model_max_budget=budget) + database.flush() + + usage: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=KEY_HASH, + model_max_budget=budget, + cache=DualCache(), + now=datetime.fromtimestamp(clock.now(), tz=timezone.utc), + shared_windows=database.replica(), + ) + assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 3, 4) + assert usage["claude-opus*"]["budget_reset_at"] == datetime.fromtimestamp( + WINDOW_EPOCH + 24 * 3600, tz=timezone.utc + ).isoformat() + + +@pytest.mark.asyncio +async def test_end_user_windows_stay_on_the_replica_without_redis(): + budget: Final = {"claude-opus*": FAMILY_CAP} + table: Final = _WindowSpendDatabase() + limiter: Final = _family_limiter(shared_windows=table.replica()) + + await _log_success( + limiter, + { + "call_type": "acompletion", + "standard_logging_object": { + "call_type": "acompletion", + "response_cost": CHAT_COST, + "model": "anthropic/claude-opus-4-8", + "model_group": "claude-opus-4-8", + "end_user": "end-user-1", + "metadata": {}, + }, + "litellm_params": {"metadata": {"user_api_key_end_user_model_max_budget": budget}}, + }, + ) + table.flush() + + assert dict(table.rows) == {} + with pytest.raises(litellm.BudgetExceededError, match="model=claude-opus-5-5"): + await limiter.is_end_user_within_model_budget( + end_user_id="end-user-1", end_user_model_max_budget=budget, model="claude-opus-5-5" + ) + + +def test_an_access_group_named_after_a_served_model_never_widens_that_models_entry(): + router: Final = Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"}, + "model_info": {"access_groups": ["gpt-4o", "openai-models"]}, + }, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": {"access_groups": ["gpt-4o", "openai-models"]}, + }, + ] + ) + + def governing_entry(model_max_budget: Mapping[str, object]) -> str | None: + resolved: Final = resolve_model_budget( + model="gpt-4o-mini", + model_max_budget=model_max_budget, + model_access_groups=lambda model, team_id: model_access_groups_of(router=router, model=model, team_id=team_id), + ) + return None if resolved is None else resolved.budget_model + + assert governing_entry({"gpt-4o": FAMILY_CAP}) is None + assert governing_entry({"openai-models": FAMILY_CAP}) == "openai-models" + assert routed_models_of(router).members("gpt-4o") == frozenset() + assert routed_models_of(router).members("openai-models") == frozenset({"gpt-4o", "gpt-4o-mini"}) + + +@pytest.mark.asyncio +async def test_key_info_usage_reports_the_family_total(): + budget: Final = {"claude-opus*": ROOMY_CAP} + limiter: Final = _family_limiter() + + await _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget) + await _spend_on(limiter, model="claude-opus-5-5", model_max_budget=budget) + + usage: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=KEY_HASH, + model_max_budget=budget, + cache=limiter.dual_cache, + ) + assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 2, 4) + + +@pytest.mark.parametrize( + ("pattern", "model", "governs"), + [ + ("claude-opus*", "claude-opus-4-8", True), + ("claude-opus*", "anthropic/claude-opus-4-8", True), + ("anthropic/claude-opus*", "anthropic/claude-opus-4-8", True), + ("anthropic/claude-opus*", "claude-opus-4-8", False), + ("claude-opus*", "my-claude-opus-4-8", False), + ("claude-opus*", "claude-sonnet-5", False), + ("gpt-4.1*", "gpt-401-mini", False), + ("*", "claude-sonnet-5", True), + ], +) +def test_glob_entries_match_the_whole_model_name_with_only_star_as_a_wildcard(pattern: str, model: str, governs: bool): + resolved: Final = resolve_model_budget( + model=model, model_max_budget={pattern: FAMILY_CAP}, model_access_groups=_model_access_groups({}) + ) + + assert (resolved is not None and resolved.budget_model == pattern) is governs + + +@pytest.mark.asyncio +async def test_replicas_without_redis_enforce_one_family_cap_together(): + budget: Final = {"claude-opus*": FAMILY_CAP} + table: Final = _WindowSpendDatabase() + replica_a: Final = _family_limiter(shared_windows=table.replica()) + replica_b: Final = _family_limiter(shared_windows=table.replica()) + + await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget) + table.flush() + + with pytest.raises(litellm.BudgetExceededError, match="model=claude-opus-5-5"): + await replica_b.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5") + + +@pytest.mark.asyncio +async def test_replicas_sharing_redis_enforce_the_family_cap_without_the_database(): + budget: Final = {"claude-opus*": FAMILY_CAP} + redis: Final = _SharedRedisDouble() + table: Final = _WindowSpendDatabase() + replica_a: Final = _family_limiter( + dual_cache=DualCache(redis_cache=redis), # pyright: ignore[reportArgumentType] # duck-typed Redis double + shared_windows=table.replica(), + ) + replica_b: Final = _family_limiter( + dual_cache=DualCache(redis_cache=redis), # pyright: ignore[reportArgumentType] # duck-typed Redis double + shared_windows=table.replica(), + ) + + await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget) + await _drain_redis_pushes() + table.flush() + + with pytest.raises(litellm.BudgetExceededError, match="model=claude-opus-5-5"): + await replica_b.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5") + assert dict(table.rows) == {} + + +@pytest.mark.asyncio +async def test_a_replica_charges_the_open_shared_window_instead_of_restarting_it(): + budget: Final = {"claude-opus*": {"budget_limit": CHAT_COST * 1.5, "time_period": "1d"}} + table: Final = _WindowSpendDatabase() + clock: Final = _Clock() + clock.advance(WINDOW_EPOCH) + replica_a: Final = _family_limiter(shared_windows=table.replica(), clock=clock.now) + replica_b: Final = _family_limiter(shared_windows=table.replica(), clock=clock.now) + + await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget) + table.flush() + clock.advance(3600) + await _spend_on(replica_b, model="claude-opus-5-5", model_max_budget=budget) + table.flush() + + for replica in (replica_a, replica_b): + with pytest.raises(litellm.BudgetExceededError): + await replica.is_key_within_model_budget(user_api_key_dict=_key(budget), model="claude-opus-5-5") + + +@pytest.mark.asyncio +async def test_key_info_on_one_replica_reports_spend_from_every_replica(): + budget: Final = {"claude-opus*": ROOMY_CAP} + table: Final = _WindowSpendDatabase() + replica_a: Final = _family_limiter(shared_windows=table.replica()) + replica_b: Final = _family_limiter(shared_windows=table.replica()) + + await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget) + await _spend_on(replica_b, model="claude-opus-5-5", model_max_budget=budget) + table.flush() + + usage: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=KEY_HASH, + model_max_budget=budget, + cache=replica_a.dual_cache, + shared_windows=table.replica(), + ) + assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 2, 4) + assert usage["claude-opus*"]["budget_reset_at"] is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("budget_model", ["claude-opus*", "claude-opus-4-8"]) +async def test_replicas_starting_a_window_moments_apart_keep_each_others_spend(budget_model: str): + budget: Final = {budget_model: ROOMY_CAP} + table: Final = _WindowSpendDatabase() + clock: Final = _Clock() + clock.advance(WINDOW_EPOCH) + replica_a: Final = _family_limiter(shared_windows=table.replica(), clock=clock.now) + replica_b: Final = _family_limiter(shared_windows=table.replica(), clock=clock.now) + + await _spend_on(replica_a, model="claude-opus-4-8", model_max_budget=budget) + clock.advance(1) + await _spend_on(replica_b, model="claude-opus-4-8", model_max_budget=budget) + table.flush() + + usage: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=KEY_HASH, + model_max_budget=budget, + cache=DualCache(), + now=datetime.fromtimestamp(clock.now(), tz=timezone.utc), + shared_windows=table.replica(), + ) + assert usage[budget_model]["current_spend"] == round(CHAT_COST * 2, 4) + + +@pytest.mark.asyncio +async def test_concurrent_first_charges_on_one_replica_both_count(): + budget: Final = {"claude-opus*": ROOMY_CAP} + clock: Final = _Clock() + clock.advance(WINDOW_EPOCH) + limiter: Final = _family_limiter(shared_windows=_WindowSpendDatabase().replica(), clock=clock.now) + + await asyncio.gather( + _spend_on(limiter, model="claude-opus-4-8", model_max_budget=budget), + _spend_on(limiter, model="claude-opus-5-5", model_max_budget=budget), + ) + + usage: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id=KEY_HASH, + model_max_budget=budget, + cache=limiter.dual_cache, + now=datetime.fromtimestamp(clock.now(), tz=timezone.utc), + shared_windows=_WindowSpendDatabase().replica(), + ) + assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 2, 4) diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 3d1e494e36c..8245c261a12 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -15,28 +15,38 @@ Redis. """ import asyncio +import base64 +import copy +import io import json from datetime import datetime, timedelta from typing import Any, Dict, Final import pytest +import soundfile as sf +from fastapi import HTTPException +from PIL import Image +from litellm import token_counter from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, estimate_image_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _AUDIO_BYTES_PER_TOKEN, + DEFAULT_AUDIO_TOKEN_ESTIMATE, PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, RateLimitedModel, - _AUDIO_BYTES_PER_TOKEN, - PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, -) -from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _call_id_from_callback_kwargs, _request_stash, get_or_create_request_stash, get_request_stash, ) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, +) from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.llms.openai import ( InputTokensDetails, ResponseAPIUsage, @@ -3775,3 +3785,284 @@ def test_reservation_counts_a_compaction_item_by_its_summary(rate_limiter): ) == handler._estimate_input_and_output_tokens( data={"model": "gpt-4o", "input": _COMPACTED_INPUT[2:], "max_output_tokens": 1}, call_type="aresponses" ) + + +_QUOTA_TEXT: Final = "a" * 400 +_QUOTA_JSON_TEXT: Final = '{"type":"input_image","image_url":"data:image/png;base64,AAAA"}' + + +@pytest.mark.parametrize( + ("call_type", "request_data", "expected"), + [ + ("acompletion", {"messages": [{"role": "user", "content": _QUOTA_TEXT}], "max_tokens": 16}, (100, 16)), + ("acompletion", {"messages": [{"role": "user", "content": _QUOTA_JSON_TEXT}], "max_tokens": 16}, (15, 16)), + ( + "anthropic_messages", + { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "t", + "content": [{"type": "text", "text": _QUOTA_TEXT}], + } + ], + } + ], + "max_tokens": 16, + }, + (0, 16), + ), + ( + "acompletion", + { + "messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "t", "type": "function", "function": {"name": "f", "arguments": _QUOTA_TEXT}} + ], + } + ], + "max_tokens": 16, + }, + (0, 16), + ), + ("aresponses", {"input": _QUOTA_TEXT, "max_output_tokens": 16}, (100, 16)), + ("aresponses", {"input": _QUOTA_TEXT, "instructions": "b" * 400, "max_output_tokens": 16}, (100, 16)), + ( + "aresponses", + {"input": [{"role": "user", "content": _QUOTA_TEXT}], "instructions": "b" * 400, "max_output_tokens": 16}, + (107, 16), + ), + ( + "aresponses", + { + "input": [ + { + "type": "file_search_call", + "id": "f", + "status": "completed", + "queries": ["q"], + "results": [ + { + "file_id": "file_1", + "filename": "result.txt", + "attributes": {}, + "score": 1, + "text": _QUOTA_TEXT, + } + ], + } + ], + "max_output_tokens": 16, + }, + (146, 16), + ), + ( + "aresponses", + { + "input": [{"type": "custom_tool_call_output", "call_id": "t", "output": _QUOTA_JSON_TEXT}], + "max_output_tokens": 16, + }, + (32, 16), + ), + ("aresponses", {"input": "", "instructions": "b" * 400}, (0, 0)), + ("aembedding", {"input": [_QUOTA_TEXT]}, (100, 0)), + ("atext_completion", {"prompt": [_QUOTA_TEXT], "max_tokens": 16}, (100, 16)), + ], +) +def test_media_correction_preserves_exact_non_media_reservations( + rate_limiter: tuple[RateLimitHandler, DualCache], + call_type: str, + request_data: dict[str, object], + expected: tuple[int, int], +) -> None: + handler, _ = rate_limiter + assert handler._estimate_input_and_output_tokens(request_data, call_type=call_type) == expected + + +def _quota_media_payload(kind: str, compressed: bool) -> str: + output: Final = io.BytesIO() + if kind == "image": + Image.new("RGB", (1024, 1024)).save(output, format="PNG", compress_level=9 if compressed else 0) + else: + sample_rate: Final = 8000 if compressed else 48000 + with sf.SoundFile(output, "w", samplerate=sample_rate, channels=1, format="WAV") as audio: + audio.buffer_write(b"\0" * (sample_rate * 2), dtype="int16") + return base64.b64encode(output.getvalue()).decode() + + +def _quota_media_request(kind: str, payload: str, call_type: str) -> dict[str, object]: + block: Final = ( + {"type": "input_audio", "input_audio": {"data": payload, "format": "wav"}} + if kind == "audio" + else {"type": "input_image", "image_url": "data:image/png;base64," + payload} + if call_type == "aresponses" + else {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": payload}} + if call_type == "anthropic_messages" + else {"type": "image_url", "image_url": {"url": "data:image/png;base64," + payload}} + ) + content: Final = ( + [{"type": "tool_result", "tool_use_id": "t", "content": [block]}] + if call_type == "anthropic_messages" + else [block] + ) + field: Final = "input" if call_type == "aresponses" else "messages" + return { + "model": "gpt-6-astra", + field: [{"role": "user", "content": content}], + "max_tokens": 16, + "max_output_tokens": 16, + } + + +@pytest.mark.parametrize( + ("kind", "call_type"), + [ + ("image", "acompletion"), + ("image", "aresponses"), + ("image", "anthropic_messages"), + ("audio", "acompletion"), + ("audio", "aresponses"), + ], +) +def test_media_reservation_is_independent_of_encoding_and_preserves_request( + rate_limiter: tuple[RateLimitHandler, DualCache], kind: str, call_type: str +) -> None: + handler, _ = rate_limiter + compressed: Final = _quota_media_request(kind, _quota_media_payload(kind, True), call_type) + uncompressed: Final = _quota_media_request(kind, _quota_media_payload(kind, False), call_type) + original: Final = copy.deepcopy(uncompressed) + estimate: Final = handler._estimate_input_and_output_tokens(uncompressed, call_type=call_type) + assert estimate == handler._estimate_input_and_output_tokens(compressed, call_type=call_type) + assert estimate[0] > 0 and estimate[1] == 16 + assert uncompressed == original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "aresponses"]) +@pytest.mark.parametrize("limit", [1000, 20000]) +async def test_media_pre_call_rejects_tight_quota_and_reserves_sufficient_quota( + rate_limiter: tuple[RateLimitHandler, DualCache], call_type: str, limit: int +) -> None: + handler, cache = rate_limiter + request: Final = { + key: value + for key, value in _quota_media_request("image", _quota_media_payload("image", True), call_type).items() + if key not in ("max_tokens", "max_output_tokens") + } + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-media-quota"), tpm_limit=limit) + if limit == 1000: + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type=call_type) + assert error.value.status_code == 429 + return + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type=call_type) + stash: Final = get_request_stash() + assert stash is not None + expected: Final = handler._estimate_input_and_output_tokens( + request, call_type=call_type, min_configured_tpm_limit=limit + ) + assert stash.reserved_tokens == sum(expected) + assert expected[0] > 1000 and expected[1] >= expected[0] + + +@pytest.mark.asyncio +async def test_ptu_pre_call_weights_media_input_and_output() -> None: + cache: Final = DualCache() + ceiling: Final = PTUTeamCeiling( + model_group="gpt-6-astra", tpm_limit=10000, output_to_input_ratio=4, cached_input_ratio=0 + ) + + def resolve_ceiling(team: str, model: str) -> PTUTeamCeiling | None: + return ceiling if (team, model) == ("test-team", "gpt-6-astra") else None + + handler: Final = RateLimitHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve_ceiling + ) + request: Final = _quota_media_request("image", _quota_media_payload("image", True), "acompletion") + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-media-ptu"), team_id="test-team"), + cache=cache, + data=request, + call_type="acompletion", + ) + stash: Final = get_request_stash() + assert stash is not None + assert ( + stash.ptu_reserved_tokens + == estimate_image_tokens("gpt-6-astra", "data:image/png;base64," + _quota_media_payload("image", True)) + 4 * 16 + ) + + +@pytest.mark.asyncio +async def test_project_fallback_keeps_legacy_audio_estimate_and_output_policy( + rate_limiter: tuple[RateLimitHandler, DualCache], +) -> None: + handler, cache = rate_limiter + unsupported: Final = {"type": "unsupported_content"} + with pytest.raises(UnsupportedContentTypeError): + token_counter(model="gpt-6-astra", messages=[{"role": "user", "content": [unsupported]}]) + request: Final = { + "model": "gpt-6-astra", + "messages": [ + { + "role": "user", + "content": [ + unsupported, + { + "type": "input_audio", + "input_audio": {"data": _quota_media_payload("audio", True), "format": "wav"}, + }, + ], + } + ], + } + auth: Final = UserAPIKeyAuth( + api_key=hash_token("sk-media-project"), + project_id="media-project", + project_metadata={"model_itpm_limit": {"gpt-6-astra": 1000}, "model_otpm_limit": {"gpt-6-astra": 1000}}, + ) + await handler.async_pre_call_hook(user_api_key_dict=auth, cache=cache, data=request, call_type="acompletion") + stash: Final = get_request_stash() + assert stash is not None + assert stash.itpm_reserved_tokens == DEFAULT_AUDIO_TOKEN_ESTIMATE + assert stash.otpm_reserved_tokens == 1 + + +@pytest.mark.asyncio +async def test_single_anthropic_image_exceeds_tight_quota_with_explicit_output_cap( + rate_limiter: tuple[RateLimitHandler, DualCache], +) -> None: + handler, cache = rate_limiter + request: Final = { + "model": "anthropic/claude-haiku-5-5", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": _quota_media_payload("image", True), + }, + } + ], + } + ], + "max_tokens": 64, + } + with pytest.raises(HTTPException) as error: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-claude-media-quota"), tpm_limit=500), + cache=cache, + data=request, + call_type="anthropic_messages", + ) + assert error.value.status_code == 429 diff --git a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py index 748b291d5ff..c13e3d5812b 100644 --- a/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_auto_router_endpoints.py @@ -170,7 +170,7 @@ async def test_preview_and_serving_share_default_model_resolution( explicit_default: str | None, expected: str, ): - from litellm.router_utils.auto_router_model_naming import validate_complexity_router_config_write + from litellm.router_strategy.auto_router.auto_router_model_naming import validate_complexity_router_config_write from litellm.types.management_endpoints.auto_router_endpoints import ComplexityRouterConfigValidationRequest config: Final = { @@ -1240,7 +1240,7 @@ class TestAutoRouterBenchmarks: def test_the_listed_kinds_match_the_router_types_traffic_can_record(self): from typing import get_args, get_type_hints - from litellm.router_utils.auto_router_model_naming import StrategyRouterKind + from litellm.router_strategy.auto_router.auto_router_model_naming import StrategyRouterKind from litellm.types.utils import StandardLoggingRoutingDecision readonly_router_type: Final = get_type_hints(StandardLoggingRoutingDecision, include_extras=True)["router_type"] diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index 8fbead8c01c..2a273034a97 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -248,6 +248,7 @@ class _FakeRepository: def __init__(self, rows: tuple[_Activity, ...]) -> None: self._rows: Final = rows self.aggregated = AsyncMock(side_effect=self._aggregated) + self.guardrail_spend = AsyncMock(return_value=()) self.key_page = AsyncMock(side_effect=self._key_page) self.key_page_call: tuple[DailyActivityScope, int, int] | None = None self.search_keys = AsyncMock(side_effect=self._search_keys) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 9634f1efce8..b13a7bfc33e 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -16580,6 +16580,7 @@ async def test_info_key_fn_includes_model_max_budget_usage(monkeypatch): } mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() monkeypatch.setattr( @@ -16703,6 +16704,7 @@ async def test_info_key_fn_v2_includes_model_max_budget_usage(monkeypatch): model_max_budget = {"gpt-4o": {"budget_limit": 1.00, "time_period": "7d"}} mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() monkeypatch.setattr( @@ -16767,6 +16769,7 @@ async def test_info_key_fn_budget_table_fallback(monkeypatch): } mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() monkeypatch.setattr( @@ -16840,6 +16843,7 @@ async def test_info_key_fn_v2_budget_table_fallback(monkeypatch): } mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() monkeypatch.setattr( @@ -17215,6 +17219,7 @@ async def test_info_key_fn_reads_the_configured_budget_model_key(monkeypatch): } mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) mock_user_api_key_cache = AsyncMock() monkeypatch.setattr( diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index fc6749aed0f..437d948f8f5 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -1,6 +1,6 @@ -import inspect import asyncio import contextlib +import inspect import json import socket from collections.abc import Iterator, Mapping @@ -1461,10 +1461,10 @@ class TestDeleteModelClearsRouterRegistry: add_deployment never restores, so an unguarded pop would make it permanently unroutable (the same cross-tenant DoS clear_cache was hardened against). """ + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) - from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete model_id = "regular-del-1" admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) @@ -2739,11 +2739,11 @@ class TestAddAndDeleteModelLifecycle: - Delete same model again → raises (model not found) """ from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, add_new_model, - delete_model as delete_model_endpoint, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( - ModelInfoDelete, + delete_model as delete_model_endpoint, ) model_id = "lifecycle-test-model-123" @@ -2833,6 +2833,8 @@ class TestDeleteTeamBYOKModelGhost: async def test_delete_strips_public_name_and_refreshes_cache(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -2916,6 +2918,8 @@ class TestDeleteTeamBYOKModelGhost: alias cleanup (delete_team_model_alias), preserving legacy behavior.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -2984,6 +2988,8 @@ class TestDeleteTeamBYOKModelGhost: team.models when one replica is deleted but a sibling still backs it.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3065,6 +3071,8 @@ class TestDeleteTeamBYOKModelGhost: deployment still serves it.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3145,6 +3153,8 @@ class TestDeleteTeamBYOKModelGhost: the alias would break routing that works.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3252,6 +3262,8 @@ class TestDeleteModelTeamAuth: async def test_proxy_admin_can_delete_model_when_team_deleted(self): from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) @@ -3287,6 +3299,8 @@ class TestDeleteModelTeamAuth: """A missing team must never let a non-admin delete the orphan (no fail-open).""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) from litellm.proxy.proxy_server import ProxyException @@ -3322,6 +3336,8 @@ class TestDeleteModelTeamAuth: """The auth check must not add a redundant team query on the live-team path.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) from litellm.proxy.proxy_server import ProxyException @@ -4392,7 +4408,6 @@ class TestModelInfoCostMapEchoFilter: def test_echoed_cost_map_metadata_is_not_persisted(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4418,7 +4433,6 @@ class TestModelInfoCostMapEchoFilter: def test_an_edited_value_survives_the_echo_filter(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4483,7 +4497,6 @@ class TestModelInfoCostMapEchoFilter: def test_metadata_without_a_cost_map_key_is_persisted(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from litellm.types.utils import echoed_cost_map_fields @@ -4513,7 +4526,6 @@ class TestModelInfoCostMapEchoFilter: def test_a_stored_mode_survives_an_echoed_save(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4536,7 +4548,6 @@ class TestModelInfoCostMapEchoFilter: def test_resetting_an_override_to_the_cost_map_value_removes_it(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4560,7 +4571,6 @@ class TestModelInfoCostMapEchoFilter: def test_reset_is_recognised_after_the_router_registered_the_override(self, monkeypatch: pytest.MonkeyPatch): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4586,7 +4596,6 @@ class TestModelInfoCostMapEchoFilter: from types import MappingProxyType import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4611,7 +4620,6 @@ class TestModelInfoCostMapEchoFilter: def test_echo_is_compared_against_the_deployments_lookup_not_the_key(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4653,7 +4661,6 @@ class TestModelInfoCostMapEchoFilter: def test_base_model_wins_over_litellm_params_model_for_the_lookup(self): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -4681,7 +4688,6 @@ class TestModelInfoCostMapEchoFilter: def test_encrypted_stored_model_is_decrypted_for_the_lookup(self, monkeypatch): import litellm - from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo @@ -6431,7 +6437,7 @@ class TestStrategyRouterWriteValidation: AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY, _auto_router_capability_slot, ) - from litellm.router_utils.auto_router_model_naming import gated_capabilities_of + from litellm.router_strategy.auto_router.auto_router_model_naming import gated_capabilities_of capabilities: Final = gated_capabilities_of(effective_params) @@ -6568,7 +6574,7 @@ class TestStrategyRouterWriteValidation: from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines configs = { "_TUNED_A": self._TUNED_A, @@ -6624,7 +6630,7 @@ class TestStrategyRouterWriteValidation: from litellm.proxy import proxy_server from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines fake = self._FakeDb([]) fake.tx_obj.litellm_proxymodeltable.find_many = AsyncMock( @@ -6669,7 +6675,7 @@ class TestStrategyRouterWriteValidation: row id before the slot is entered, so a second tuned router is refused before its DB write.""" from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model - from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) baselines = snapshot_tuning_baselines([self._db_router_row("a", self._TUNED_A)]) @@ -8980,7 +8986,10 @@ class TestTeamMemberAutoRouterWrites: from fastapi import HTTPException from litellm.proxy.management_endpoints.model_management_endpoints import _auto_router_capability_slot - from litellm.proxy.management_helpers.auto_router_permissions import MemberAutoRouterWrite, validate_member_auto_router_config + from litellm.proxy.management_helpers.auto_router_permissions import ( + MemberAutoRouterWrite, + validate_member_auto_router_config, + ) row: Final = self._row() database: Final = self._database(self._team(), row) @@ -10677,6 +10686,8 @@ class TestFederationGateScopesToWhatTheWriteTouches: and the team that owns the model must be able to take it off their page.""" from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index f727448f3ac..e51f26bf1c3 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -44,6 +44,7 @@ from litellm.proxy._types import ( UpdateTeamRequest, UserAPIKeyAuth, # Import UserAPIKeyAuth ) +from litellm.proxy.hooks.model_max_budget_limiter import RoutedModels, routed_models_of from litellm.proxy.management.teams.authz import TeamAccess from litellm.proxy.management_endpoints.team_admin_field_permissions import ( PROXY_CONTROLLED_TEAM_METADATA_KEYS, @@ -17570,6 +17571,26 @@ def test_team_admin_cannot_loosen_team_model_caps(requested) -> None: assert "proxy admin" in exc.value.detail["error"] +def test_team_cap_refusal_names_the_entry_holding_the_cap() -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + with pytest.raises(HTTPException) as exc: + _check_team_model_budget_update_authority( + data=UpdateTeamRequest( + team_id="t1", + model_max_budget={ + **_EXISTING_TEAM_MODEL_CAPS, + "openai/gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}, + }, + ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ) + assert exc.value.status_code == 403 + assert "for 'openai/gpt-4o'" in exc.value.detail["error"] + assert "(entry 'gpt-4o')" in exc.value.detail["error"] + + @pytest.mark.parametrize( "requested", [ @@ -17593,6 +17614,253 @@ def test_team_admin_can_tighten_or_keep_team_model_caps(requested) -> None: ) +_TEAM_OPUS_FAMILY_CAP: Final = MappingProxyType({"claude-opus*": {"max_budget": 10.0, "budget_duration": "1d"}}) +_TEAM_OPUS_GROUP_CAP: Final = MappingProxyType({"z-opus-group": {"max_budget": 10.0, "budget_duration": "1d"}}) +_TEAM_ROUTED_MODELS: Final = RoutedModels( + served_names=frozenset( + {"claude-opus-4-8", "claude-opus-5-5", "gemini-3-pro", "gpt-4o", "my-team/gpt-4o", "team-route/*"} + ), + group_members=MappingProxyType( + { + "opus-group": frozenset({"claude-opus-4-8"}), + "z-opus-group": frozenset({"claude-opus-4-8"}), + "a-opus-group": frozenset({"claude-opus-4-8", "claude-opus-5-5"}), + "zz-opus-group": frozenset({"claude-opus-4-8"}), + "a-gemini-group": frozenset({"gemini-3-pro"}), + "z-anthropic-group": frozenset({"anthropic/*"}), + } + ), +) + + +def _first_bedrock_anthropic_model(model_cost: Mapping[str, Mapping[str, object]]) -> str: + return next( + model.removeprefix("anthropic.") + for model, cost in model_cost.items() + if model.startswith("anthropic.") and str(cost.get("litellm_provider", "")).startswith("bedrock") + ) + + +_BEDROCK_ANTHROPIC_MODEL: Final = _first_bedrock_anthropic_model(litellm.model_cost) +_TEAM_BEDROCK_VENDOR_CAP: Final = MappingProxyType({"anthropic.*": {"max_budget": 10.0, "budget_duration": "1d"}}) +_BEDROCK_ROUTE_ADDED_ENTRIES: Final = (_BEDROCK_ANTHROPIC_MODEL, f"{_BEDROCK_ANTHROPIC_MODEL[:12]}*") + + +def _check_team_admin_model_caps( + existing: Mapping[str, object], + requested: Mapping[str, object], + routed_models: RoutedModels = _TEAM_ROUTED_MODELS, +) -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=dict(requested)), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=existing, + routed_models_source=lambda _team_id: routed_models, + ) + + +@pytest.mark.parametrize( + ("existing", "added"), + [ + (_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 10.0, "budget_duration": "1h"}}), + (_TEAM_OPUS_FAMILY_CAP, {"*laude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {"a-opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}), + ( + {"z-anthropic-group": {"max_budget": 10.0, "budget_duration": "1d"}}, + {"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ), + (_TEAM_OPUS_FAMILY_CAP, {"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}), + ( + {"anthropic/*": {"max_budget": 10.0, "budget_duration": "1d"}}, + {"claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ), + ( + {"my-team/*": {"max_budget": 10.0, "budget_duration": "1d"}}, + {"gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ), + ( + {"team-route/*": {"max_budget": 10.0, "budget_duration": "1d"}}, + {"gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ), + (_TEAM_OPUS_FAMILY_CAP, {"claude-opus-4-8": {"max_budget": 10.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_FAMILY_CAP, {"claude-opus-4*": {"max_budget": 10.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 5.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 10.0, "budget_duration": "2d"}}), + (_TEAM_OPUS_FAMILY_CAP, {"opus-group": {"max_budget": 0.01, "budget_duration": "1h"}}), + (_TEAM_OPUS_FAMILY_CAP, {"claude-opus-4-8": {"rpm_limit": 10}}), + ], + ids=[ + "group_over_pattern", + "group_with_a_faster_window", + "more_specific_pattern", + "earlier_group_sharing_a_model", + "model_inside_group", + "model_under_wildcard_group", + "model_matching_pattern", + "model_under_provider_pattern", + "model_under_served_prefix", + "model_under_wildcard_route", + "model_carved_out_at_the_same_cap", + "pattern_carved_out_at_the_same_cap", + "tighter_group", + "group_with_a_slower_window", + "group_with_a_faster_window_within_the_cap", + "uncapped_model", + ], +) +def test_team_admin_cannot_add_an_entry_that_takes_precedence_over_a_cap( + existing: Mapping[str, object], added: Mapping[str, object] +) -> None: + with pytest.raises(HTTPException) as exc: + _check_team_admin_model_caps(existing=existing, requested={**existing, **added}) + assert exc.value.status_code == 403 + assert "proxy admin" in exc.value.detail["error"] + assert repr(next(iter(existing))) in exc.value.detail["error"] + + +@pytest.mark.parametrize( + ("existing", "requested"), + [ + (_TEAM_OPUS_FAMILY_CAP, {**_TEAM_OPUS_FAMILY_CAP, "claude*": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "a-gemini-group": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "zz-opus-group": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_GROUP_CAP, {**_TEAM_OPUS_GROUP_CAP, "gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}), + (_TEAM_OPUS_FAMILY_CAP, {**_TEAM_OPUS_FAMILY_CAP, "gpt-4o": {"max_budget": 1000.0, "budget_duration": "1d"}}), + ( + _TEAM_OPUS_FAMILY_CAP, + {**_TEAM_OPUS_FAMILY_CAP, "a-gemini-group": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ), + ( + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}}, + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}}, + ), + ( + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 100.0, "budget_duration": "1d"}}, + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"max_budget": 50.0, "budget_duration": "1d"}}, + ), + ( + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"rpm_limit": 10}}, + {**_TEAM_OPUS_FAMILY_CAP, "opus-group": {"rpm_limit": 20}}, + ), + ( + {**_TEAM_OPUS_FAMILY_CAP, "claude-opus-4-8": {"rpm_limit": 10, "budget_duration": "1d"}}, + {**_TEAM_OPUS_FAMILY_CAP, "claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ), + ( + _TEAM_OPUS_FAMILY_CAP, + {**_TEAM_OPUS_FAMILY_CAP, "gemini-3-pro*": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ), + ( + {**_TEAM_OPUS_FAMILY_CAP, "gpt-4o": {"max_budget": 10.0}}, + {**_TEAM_OPUS_FAMILY_CAP, "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}}, + ), + ( + {"z-anthropic-group": {"max_budget": 10.0}}, + {"z-anthropic-group": {"max_budget": 10.0, "budget_duration": "1d"}}, + ), + ], + ids=[ + "less_specific_pattern", + "group_without_shared_models", + "later_group_sharing_a_model", + "model_outside_group", + "model_outside_pattern", + "group_outside_pattern", + "unchanged_looser_entry", + "tightened_looser_entry", + "uncapped_entry_edited", + "enforced_uncapped_entry_capped", + "pattern_sharing_only_a_wildcard_route", + "unenforced_capped_model_given_a_duration", + "unenforced_capped_wildcard_group_given_a_duration", + ], +) +def test_team_admin_can_keep_or_tighten_entries_and_add_ones_that_outrank_no_cap( + existing: Mapping[str, object], requested: Mapping[str, object] +) -> None: + assert _check_team_admin_model_caps(existing=existing, requested=requested) is None + + +def test_team_admin_cannot_make_an_unenforced_entry_take_precedence_over_a_cap() -> None: + existing: Final = {**_TEAM_OPUS_FAMILY_CAP, "claude-opus-4-8": {"rpm_limit": 10}} + + with pytest.raises(HTTPException) as exc: + _check_team_admin_model_caps( + existing=existing, + requested={**existing, "claude-opus-4-8": {"max_budget": 1000.0, "budget_duration": "1d"}}, + ) + assert exc.value.status_code == 403 + assert "'claude-opus*'" in exc.value.detail["error"] + + +def test_team_admin_cannot_raise_an_unenforced_cap_while_enforcing_it() -> None: + with pytest.raises(HTTPException) as exc: + _check_team_admin_model_caps( + existing={"gpt-4o": {"max_budget": 10.0}}, + requested={"gpt-4o": {"max_budget": 20.0, "budget_duration": "1d"}}, + ) + assert exc.value.status_code == 403 + assert "(entry 'gpt-4o')" in exc.value.detail["error"] + + +@pytest.mark.parametrize("added_entry", _BEDROCK_ROUTE_ADDED_ENTRIES, ids=["model", "pattern"]) +def test_team_admin_cannot_lift_a_bedrock_vendor_cap_through_a_bedrock_wildcard_route(added_entry: str) -> None: + with pytest.raises(HTTPException) as exc: + _check_team_admin_model_caps( + existing=_TEAM_BEDROCK_VENDOR_CAP, + requested={**_TEAM_BEDROCK_VENDOR_CAP, added_entry: {"max_budget": 1000.0, "budget_duration": "1d"}}, + routed_models=RoutedModels( + served_names=frozenset({"bedrock/*", "gpt-4o"}), group_members=MappingProxyType({}) + ), + ) + assert exc.value.status_code == 403 + assert "'anthropic.*'" in exc.value.detail["error"] + + +@pytest.mark.parametrize("added_entry", _BEDROCK_ROUTE_ADDED_ENTRIES, ids=["model", "pattern"]) +def test_team_admin_can_add_entries_beside_a_bedrock_vendor_cap_no_route_serves(added_entry: str) -> None: + assert ( + _check_team_admin_model_caps( + existing=_TEAM_BEDROCK_VENDOR_CAP, + requested={**_TEAM_BEDROCK_VENDOR_CAP, added_entry: {"max_budget": 1000.0, "budget_duration": "1d"}}, + routed_models=RoutedModels(served_names=frozenset({"gpt-4o"}), group_members=MappingProxyType({})), + ) + is None + ) + + +def test_team_admin_cannot_outrank_a_group_cap_through_a_team_model_both_groups_serve() -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + team_router: Final = Router( + model_list=[ + { + "model_name": f"model_name_t1_{group}", + "litellm_params": {"model": "anthropic/claude-opus-4-8", "api_key": "sk-test"}, + "model_info": {"team_id": "t1", "team_public_model_name": "claude-opus-4-8", "access_groups": [group]}, + } + for group in ("a-opus", "z-opus") + ] + ) + existing: Final = {"z-opus": {"max_budget": 10.0, "budget_duration": "1d"}} + + with pytest.raises(HTTPException) as exc: + _check_team_model_budget_update_authority( + data=UpdateTeamRequest( + team_id="t1", model_max_budget={**existing, "a-opus": {"max_budget": 100.0, "budget_duration": "1d"}} + ), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=existing, + routed_models_source=lambda team_id: routed_models_of(router=team_router, team_id=team_id), + ) + assert exc.value.status_code == 403 + assert "'z-opus'" in exc.value.detail["error"] + + def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_admins() -> None: from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority diff --git a/tests/unit/proxy/management_helpers/test_auto_router_availability.py b/tests/unit/proxy/management_helpers/test_auto_router_availability.py index a88dcd7f66d..0eaae60b4ea 100644 --- a/tests/unit/proxy/management_helpers/test_auto_router_availability.py +++ b/tests/unit/proxy/management_helpers/test_auto_router_availability.py @@ -1,16 +1,15 @@ from collections.abc import Mapping -from typing import Final from types import SimpleNamespace - -from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from typing import Final import pytest +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.management_helpers.auto_router_availability import ( auto_router_availability, build_auto_router_catalog, ) -from litellm.router_utils.auto_router_tuning_baseline import snapshot_tuning_baselines +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import snapshot_tuning_baselines def deployment( diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index d47db69257b..ff63e37bcd8 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -1,3 +1,4 @@ +import itertools import json from typing import Final, List, Optional from urllib.parse import quote @@ -12,6 +13,7 @@ from pytest_mock import MockerFixture import litellm from litellm import Router +from litellm.constants import ANTHROPIC_DEFAULT_FILE_LIST_LIMIT from litellm.files.types import FileContentStreamingResult from litellm.proxy._types import LiteLLM_UserTableFiltered, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.hooks import get_proxy_hook @@ -4969,6 +4971,8 @@ def _managed_file_row( stored_file_id: Optional[str] = None, created_by: str = "test-user", team_id: Optional[str] = None, + model_mappings: Optional[dict[str, str]] = None, + storage_backend: Optional[str] = None, ) -> MagicMock: file_object = OpenAIFileObject( id=stored_file_id or unified_file_id, @@ -4984,6 +4988,8 @@ def _managed_file_row( file_object=file_object.model_dump(), created_by=created_by, team_id=team_id, + model_mappings=model_mappings or {}, + storage_backend=storage_backend, ) @@ -5010,16 +5016,49 @@ class _ManagedFileTableOverRows: return next(iter(self._owned_rows(where)), None) async def find_many(self, where, take=None, order=None, cursor=None, skip=0): - rows = self._owned_rows(where) - if cursor is not None: - start = next( - index - for index, row in enumerate(rows) - if row.unified_file_id == cursor["unified_file_id"] - ) - rows = rows[start + skip :] + owned_rows: Final = self._owned_rows(where) + oldest_first: Final = bool(order) and order[0].get("created_at") == "asc" + ordered_rows: Final = owned_rows[::-1] if oldest_first else owned_rows + start: Final = ( + 0 + if cursor is None + else skip + + next(index for index, row in enumerate(ordered_rows) if row.unified_file_id == cursor["unified_file_id"]) + ) + rows: Final = ordered_rows[start:] return rows if take is None else rows[:take] + async def query_raw(self, query: str, *params: object) -> list[dict[str, object]]: + keys_index: Final = next(index for index, param in enumerate(params) if isinstance(param, list)) + owner_ids: Final = frozenset(params[:keys_index]) + model_mapping_keys: Final = frozenset(params[keys_index]) + cursor_id: Final = params[keys_index + 1] if len(params) == keys_index + 3 else None + ordered_rows: Final = [ + row + for row in (self.rows[::-1] if '"created_at" ASC' in query else self.rows) + if not owner_ids or {row.created_by, row.team_id} & owner_ids + ] + start: Final = ( + 0 + if cursor_id is None + else 1 + next(index for index, row in enumerate(ordered_rows) if row.unified_file_id == cursor_id) + ) + listed_rows: Final = [ + row + for row in ordered_rows[start:] + if not model_mapping_keys.isdisjoint(row.model_mappings) + or ('"storage_backend" IS NOT NULL' in query and row.storage_backend is not None) + ] + return [ + { + "unified_file_id": row.unified_file_id, + "file_object": row.file_object, + "model_mappings": row.model_mappings, + "storage_backend": row.storage_backend, + } + for row in listed_rows[: params[-1]] + ] + def _setup_unscoped_list_files_route_over_real_hook( mocker, monkeypatch, llm_router: Router, rows @@ -5041,7 +5080,9 @@ def _setup_unscoped_list_files_route_over_real_hook( managed_files = PROXY_LiteLLMManagedFiles( internal_usage_cache=MagicMock(), prisma_client=MagicMock() ) - managed_files.prisma_client.db.litellm_managedfiletable = _ManagedFileTableOverRows(rows) + managed_file_table = _ManagedFileTableOverRows(rows) + managed_files.prisma_client.db.litellm_managedfiletable = managed_file_table + managed_files.prisma_client.db.query_raw = managed_file_table.query_raw proxy_logging_obj = setup_proxy_logging_object(monkeypatch, llm_router) proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files @@ -6263,3 +6304,513 @@ def test_create_file_passthrough_fails_closed_when_guardrails_would_scan_the_bat assert error["param"] == "passthrough" assert "guardrails" in error["message"] assert forwarded_calls == [] + + +def _callers_files_on_every_deployment() -> list[MagicMock]: + return [ + _managed_file_row("unified-gpt-newer", model_mappings={"gpt-3.5-turbo-id": "file-gpt-newer"}), + _managed_file_row("unified-azure", model_mappings={"azure-gpt-3-5-turbo-id": "file-azure"}), + _managed_file_row("unified-gemini", model_mappings={"gemini-2.0-flash-id": "file-gemini"}), + _managed_file_row( + "unified-theirs", created_by="other-user", model_mappings={"gpt-3.5-turbo-id": "file-theirs"} + ), + _managed_file_row("unified-gpt-older", model_mappings={"gpt-3.5-turbo-id": "file-gpt-older"}), + ] + + +def _list_files_as_test_user(path: str, params: dict[str, str], headers: dict[str, str]) -> httpx.Response: + import litellm.proxy.proxy_server as ps + + try: + return client.get(path, params=params, headers={"Authorization": "Bearer test-key", **headers}) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + +@pytest.mark.parametrize( + ("path", "params", "headers", "expected_ids"), + [ + ("/openai/v1/files", {}, {}, ["unified-gpt-newer", "unified-gpt-older"]), + ("/v1/files", {"provider": "azure"}, {}, ["unified-azure"]), + ("/v1/files", {}, {"custom-llm-provider": "gemini"}, ["unified-gemini"]), + ("/v1/files", {"target_model_names": "gemini-2.0-flash"}, {}, ["unified-gemini"]), + ("/v1/files", {}, {"x-litellm-model": "gpt-3.5-turbo"}, ["unified-gpt-newer", "unified-gpt-older"]), + ("/v1/files", {}, {}, ["unified-gpt-newer", "unified-azure", "unified-gemini", "unified-gpt-older"]), + ], + ids=["provider-path", "provider-query", "provider-header", "target-model-names", "model-header", "no-provider"], +) +def test_require_managed_files_lists_the_callers_files_on_the_named_provider_or_model( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + path: str, + params: dict[str, str], + headers: dict[str, str], + expected_ids: list[str], +) -> None: + _, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, params, headers) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == expected_ids + provider_list.assert_not_awaited() + + +def test_require_managed_files_lists_a_teammates_files_on_the_provider_and_no_other_teams( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + import litellm.proxy.proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + llm_router, + [ + _managed_file_row( + "unified-teammate", created_by="teammate", team_id="team-a", model_mappings={"gpt-3.5-turbo-id": "f-1"} + ), + _managed_file_row("unified-own", model_mappings={"gpt-3.5-turbo-id": "f-2"}), + _managed_file_row( + "unified-other-team", created_by="outsider", team_id="team-b", model_mappings={"gpt-3.5-turbo-id": "f-3"} + ), + ], + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", user_role=LitellmUserRoles.INTERNAL_USER, user_id="test-user", team_id="team-a" + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user("/openai/v1/files", {}, {}) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == ["unified-teammate", "unified-own"] + + +def _router_with_wildcard_and_unprefixed_anthropic_deployments() -> Router: + unprefixed_anthropic_model: Final = next(model for model in sorted(litellm.anthropic_models) if "/" not in model) + return Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"}, + "model_info": {"id": "gpt-3.5-turbo-id"}, + }, + { + "model_name": "anthropic/*", + "litellm_params": {"model": "anthropic/*", "api_key": "anthropic_api_key"}, + "model_info": {"id": "anthropic-wildcard-id"}, + }, + { + "model_name": "claude-unprefixed", + "litellm_params": {"model": unprefixed_anthropic_model, "api_key": "anthropic_api_key"}, + "model_info": {"id": "claude-unprefixed-id"}, + }, + ], + model_group_alias={"claude-alias": "claude-unprefixed"}, + routing_groups=[ + {"group_name": "claude-group", "models": ["claude-unprefixed"], "routing_strategy": "simple-shuffle"} + ], + ) + + +def _callers_files_on_wildcard_unprefixed_and_passthrough_deployments() -> list[MagicMock]: + return [ + _managed_file_row("unified-wildcard", model_mappings={"anthropic-wildcard-id": "file-wildcard"}), + _managed_file_row("unified-unprefixed", model_mappings={"claude-unprefixed-id": "file-unprefixed"}), + _managed_file_row("unified-gpt", model_mappings={"gpt-3.5-turbo-id": "file-gpt"}), + _managed_file_row( + "unified-passthrough", + model_mappings={ + "_passthrough_openai": "file-passthrough", + "_passthrough_provider_marker": "_passthrough_provider:openai", + }, + ), + _managed_file_row( + "unified-storage", + model_mappings={"claude-unprefixed": "litellm_db://unified-storage"}, + storage_backend="litellm_db", + ), + _managed_file_row( + "unified-wildcard-storage", + model_mappings={"anthropic/claude-haiku-5-5": "litellm_db://unified-wildcard-storage"}, + storage_backend="litellm_db", + ), + _managed_file_row( + "unified-alias-storage", + model_mappings={"claude-alias": "s3://bucket/unified-alias-storage"}, + storage_backend="s3", + ), + _managed_file_row( + "unified-group-storage", + model_mappings={"claude-group": "s3://bucket/unified-group-storage"}, + storage_backend="s3", + ), + ] + + +@pytest.mark.parametrize( + ("path", "params", "headers", "expected_ids"), + [ + ( + "/v1/files", + {"target_model_names": "anthropic/claude-haiku-5-5"}, + {}, + ["unified-wildcard", "unified-wildcard-storage"], + ), + ( + "/v1/files", + {}, + {"x-litellm-model": "anthropic/claude-haiku-5-5"}, + ["unified-wildcard", "unified-wildcard-storage"], + ), + ( + "/anthropic/v1/files", + {}, + {}, + [ + "unified-wildcard", + "unified-unprefixed", + "unified-storage", + "unified-wildcard-storage", + "unified-alias-storage", + "unified-group-storage", + ], + ), + ("/openai/v1/files", {}, {}, ["unified-gpt", "unified-passthrough"]), + ("/v1/files", {"target_model_names": "claude-unprefixed"}, {}, ["unified-unprefixed", "unified-storage"]), + ], + ids=[ + "wildcard-target-model-names", + "wildcard-model-header", + "unprefixed-provider", + "passthrough-provider", + "storage-backed-model", + ], +) +def test_require_managed_files_lists_the_callers_files_wherever_the_router_put_them( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + path: str, + params: dict[str, str], + headers: dict[str, str], + expected_ids: list[str], +) -> None: + _, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + _router_with_wildcard_and_unprefixed_anthropic_deployments(), + _callers_files_on_wildcard_unprefixed_and_passthrough_deployments(), + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, params, headers) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == expected_ids + provider_list.assert_not_awaited() + + +def _router_with_a_catch_all_openai_deployment() -> Router: + unprefixed_anthropic_model: Final = next(model for model in sorted(litellm.anthropic_models) if "/" not in model) + return Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "openai_api_key"}, + "model_info": {"id": "gpt-3.5-turbo-id"}, + }, + { + "model_name": "*", + "litellm_params": {"model": "openai/*", "api_key": "openai_api_key"}, + "model_info": {"id": "catch-all-id"}, + }, + { + "model_name": "claude-unprefixed", + "litellm_params": {"model": unprefixed_anthropic_model, "api_key": "anthropic_api_key"}, + "model_info": {"id": "claude-unprefixed-id"}, + }, + ] + ) + + +@pytest.mark.parametrize( + ("path", "expected_ids"), + [ + ("/openai/v1/files", ["unified-gpt", "unified-catch-all-storage"]), + ("/anthropic/v1/files", ["unified-unprefixed", "unified-storage"]), + ], + ids=["catch-all-provider", "explicit-provider"], +) +def test_require_managed_files_keeps_a_catch_all_route_from_claiming_files_on_other_providers( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + path: str, + expected_ids: list[str], +) -> None: + _, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + _router_with_a_catch_all_openai_deployment(), + [ + _managed_file_row("unified-unprefixed", model_mappings={"claude-unprefixed-id": "file-unprefixed"}), + _managed_file_row("unified-gpt", model_mappings={"gpt-3.5-turbo-id": "file-gpt"}), + _managed_file_row( + "unified-storage", + model_mappings={"claude-unprefixed": "litellm_db://unified-storage"}, + storage_backend="litellm_db", + ), + _managed_file_row( + "unified-catch-all-storage", + model_mappings={"model-only-the-catch-all-serves": "litellm_db://unified-catch-all-storage"}, + storage_backend="litellm_db", + ), + ], + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, {}, {}) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == expected_ids + provider_list.assert_not_awaited() + + +@pytest.mark.parametrize( + ("path", "params"), + [ + ("/antropic/v1/files", {}), + ("/openai/v1/files", {"after": "unified-gpt", "after_id": "unified-wildcard"}), + ("/v1/files", {"target_model_names": "gpt-3.5-turbo,claude-unprefixed"}), + ], + ids=["unknown-provider", "conflicting-after-cursors", "two-target-model-names"], +) +def test_require_managed_files_rejects_a_list_it_cannot_scope( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, path: str, params: dict[str, str] +) -> None: + _, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + _router_with_wildcard_and_unprefixed_anthropic_deployments(), + _callers_files_on_wildcard_unprefixed_and_passthrough_deployments(), + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, params, {}) + + assert response.status_code == 400, response.text + provider_list.assert_not_awaited() + + +@pytest.mark.parametrize("cursor_param", ["after", "after_id"]) +def test_require_managed_files_pages_a_provider_list_with_either_cursor_param( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router, cursor_param: str +) -> None: + import litellm.proxy.proxy_server as ps + + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + try: + first_page: Final = client.get( + "/openai/v1/files", params={"limit": "1"}, headers={"Authorization": "Bearer test-key"} + ) + second_page: Final = client.get( + "/openai/v1/files", + params={"limit": "1", cursor_param: first_page.json()["last_id"]}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert (first_page.status_code, second_page.status_code) == (200, 200), (first_page.text, second_page.text) + assert [file["id"] for file in first_page.json()["data"]] == ["unified-gpt-newer"] + assert first_page.json()["has_more"] is True + assert [file["id"] for file in second_page.json()["data"]] == ["unified-gpt-older"] + assert second_page.json()["has_more"] is False + + +def test_require_managed_files_pages_toward_newer_files_with_before_id( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + import litellm.proxy.proxy_server as ps + + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + try: + page_before_oldest: Final = client.get( + "/v1/files", + params={"limit": "2", "before_id": "unified-gpt-older"}, + headers={"Authorization": "Bearer test-key"}, + ) + newest_page: Final = client.get( + "/v1/files", + params={"limit": "2", "before_id": page_before_oldest.json()["first_id"]}, + headers={"Authorization": "Bearer test-key"}, + ) + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + + assert (page_before_oldest.status_code, newest_page.status_code) == (200, 200), ( + page_before_oldest.text, + newest_page.text, + ) + assert [file["id"] for file in page_before_oldest.json()["data"]] == ["unified-azure", "unified-gemini"] + assert page_before_oldest.json()["has_more"] is True + assert [file["id"] for file in newest_page.json()["data"]] == ["unified-gpt-newer"] + assert newest_page.json()["has_more"] is False + + +def test_managed_file_list_rejects_both_cursor_directions( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user( + "/v1/files", {"after": "unified-gpt-newer", "before_id": "unified-gpt-older"}, {} + ) + + assert response.status_code == 400, response.text + assert "not both" in response.json()["error"]["message"] + + +@pytest.mark.parametrize( + ("params", "headers"), + [({"target_model_names": "gemini-2.0-flash"}, {}), ({}, {"x-litellm-model": "gemini-2.0-flash"})], + ids=["target-model-names", "model-header"], +) +def test_require_managed_files_list_rejects_a_model_the_key_cannot_use( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + params: dict[str, str], + headers: dict[str, str], +) -> None: + import litellm.proxy.proxy_server as ps + + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="test-user", + models=["gpt-3.5-turbo"], + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user("/v1/files", params, headers) + + assert response.status_code == 403, response.text + assert "unified-gemini" not in response.text + + +def test_require_managed_files_list_fails_closed_without_the_managed_files_hook( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + proxy_logging_obj, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + proxy_logging_obj.proxy_hook_mapping.pop("managed_files") + provider_list.return_value = [_managed_file("file-on-the-shared-key")] + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user("/openai/v1/files", {}, {}) + + assert response.status_code == 500, response.text + assert "file-on-the-shared-key" not in response.text + assert "require_managed_files cannot scope the list to the caller" in response.json()["error"]["message"] + provider_list.assert_not_awaited() + + +@pytest.mark.parametrize("hook_loaded", [True, False], ids=["hook-loaded", "no-hook"]) +def test_require_managed_files_keeps_the_provider_list_for_an_admin( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + llm_router: Router, + hook_loaded: bool, +) -> None: + import litellm.proxy.proxy_server as ps + + proxy_logging_obj, provider_list = _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + if not hook_loaded: + proxy_logging_obj.proxy_hook_mapping.pop("managed_files") + provider_list.return_value = [_managed_file("file-on-the-shared-key")] + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + api_key="test-key", user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user" + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user("/anthropic/v1/files", {}, {}) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == ["file-on-the-shared-key"] + + +def test_managed_file_list_without_require_managed_files_ignores_cursor_params_openai_does_not_define( + mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch, llm_router: Router +) -> None: + _setup_unscoped_list_files_route_over_real_hook( + mocker, monkeypatch, llm_router, _callers_files_on_every_deployment() + ) + monkeypatch.setattr("litellm.require_managed_files", False) + + response: Final = _list_files_as_test_user( + "/v1/files", {"after": "unified-gpt-newer", "after_id": "unified-gemini", "before_id": "file-raw"}, {} + ) + + assert response.status_code == 200, response.text + assert [file["id"] for file in response.json()["data"]] == ["unified-azure", "unified-gemini", "unified-gpt-older"] + + +@pytest.mark.parametrize( + ("path", "headers", "expected_count", "has_more"), + [ + ("/anthropic/v1/files", {}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT, True), + ("/v1/files", {"x-litellm-model": "claude-unprefixed"}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT, True), + ("/anthropic/v1/files", {"x-litellm-model": "gpt-3.5-turbo"}, ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1, False), + ("/v1/files", {}, 2 * (ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1), False), + ], + ids=["anthropic-provider", "anthropic-model-header", "openai-model-header", "every-provider"], +) +def test_require_managed_files_pages_an_anthropic_list_at_the_anthropic_default_size( + mocker: MockerFixture, + monkeypatch: pytest.MonkeyPatch, + path: str, + headers: dict[str, str], + expected_count: int, + has_more: bool, +) -> None: + _setup_unscoped_list_files_route_over_real_hook( + mocker, + monkeypatch, + _router_with_wildcard_and_unprefixed_anthropic_deployments(), + [ + _managed_file_row( + f"unified-{deployment_id}-{index:03d}", model_mappings={deployment_id: f"file-{index:03d}"} + ) + for deployment_id, index in itertools.product( + ("claude-unprefixed-id", "gpt-3.5-turbo-id"), range(ANTHROPIC_DEFAULT_FILE_LIST_LIMIT + 1) + ) + ], + ) + monkeypatch.setattr("litellm.require_managed_files", True) + + response: Final = _list_files_as_test_user(path, {}, headers) + + assert response.status_code == 200, response.text + assert (len(response.json()["data"]), response.json()["has_more"]) == (expected_count, has_more) diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index e09a179eadf..47b6af9aeb3 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2503,6 +2503,33 @@ def test_stream_was_interrupted_skips_data_lines_that_are_not_json_objects( assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is interrupted + +def test_stream_was_interrupted_reads_a_message_delta_whose_json_holds_unicode_line_separators() -> None: + message_delta: Final = { + "type": "message_delta", + "delta": {"stop_reason": "stop_sequence", "stop_sequence": "END\u2028OF\u2029TURN"}, + "usage": {"output_tokens": 3}, + } + all_chunks: Final = [ + 'data: {"type": "content_block_stop", "index": 0}', + f"data: {json.dumps(message_delta, ensure_ascii=False)}", + ] + + assert AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks) is False + + +def test_extract_model_reads_a_message_start_whose_json_holds_unicode_line_separators() -> None: + message_start: Final = { + "type": "message_start", + "message": { + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "Prefill\u2028line\u2029two"}], + }, + } + chunk: Final = f"event: message_start\ndata: {json.dumps(message_start, ensure_ascii=False)}\n\n" + + assert AnthropicPassthroughLoggingHandler._extract_model_from_anthropic_chunks([chunk]) == "claude-sonnet-5" + @pytest.mark.asyncio async def test_anthropic_passthrough_handler(mock_httpx_response, mock_response, mock_logging_obj): """ @@ -3004,3 +3031,24 @@ def test_anthropic_passthrough_stream_bills_the_bedrock_guardrail_trace_of_the_m (entry,) = _recorded_guardrail_entries(mock_logging_obj) assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015) + + +def test_anthropic_passthrough_stream_bills_a_bedrock_guardrail_trace_whose_json_holds_unicode_line_separators( + monkeypatch: pytest.MonkeyPatch, mock_logging_obj: LiteLLMLoggingObj, all_chunks: list[str] +) -> None: + monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING) + mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + trace: Final = {"guardrail": {**_BEDROCK_INVOKE_TRACE["guardrail"], "actionReason": "Blocked\u2028by\u2029policy"}} + stop_event: Final = json.dumps({"type": "message_stop", "amazon-bedrock-trace": trace}, ensure_ascii=False) + AnthropicPassthroughLoggingHandler.handle_logging_anthropic_collected_chunks( + litellm_logging_obj=mock_logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "claude", "messages": []}, + endpoint_type=MagicMock(), + start_time=datetime(2025, 1, 1), + all_chunks=[*all_chunks, f"event: message_stop\ndata: {stop_event}\n\n"], + end_time=datetime(2025, 1, 1), + ) + (entry,) = _recorded_guardrail_entries(mock_logging_obj) + assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py index 1b83c5140ca..b0f84f2ecbe 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py @@ -1,6 +1,7 @@ import asyncio import json from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock import httpx @@ -267,6 +268,12 @@ class TestRunIdFromSseFrames: frames = b'data: {"run_id": "run-7", "event": "INITIALIZED"}\n\ndata: {"run_id": "run-7", "event": "ACTION"}\n\n' assert run_id_from_sse_frames(frames) == "run-7" + def test_finds_run_id_in_a_frame_whose_json_holds_unicode_line_separators(self) -> None: + event: Final = {"run_id": "run-7", "event": "ACTION", "message": "Clicked\u2028Submit\u2029button"} + frames: Final = f"data: {json.dumps(event, ensure_ascii=False)}\n\n".encode() + + assert run_id_from_sse_frames(frames) == "run-7" + def test_skips_frames_without_run_id(self): frames = b': keepalive\n\ndata: not-json\n\ndata: {"event": "HEARTBEAT"}\n\n' assert run_id_from_sse_frames(frames) is None diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py index 49bc2e9c29e..2496aa47d22 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -30,7 +30,7 @@ from litellm.llms.base_llm.managed_resources.utils import( resolve_passthrough_managed_id_provider, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome -from typing import Any +from typing import Any, Final def _user() -> UserAPIKeyAuth: @@ -254,6 +254,37 @@ async def test_streamed_response_with_cr_only_frame_delimiters_is_still_owned_an assert output == _response_stream_bytes(managed_id).replace(b"\n", b"\r") +def _instructed_response_stream_bytes(raw_id: str, instructions: str) -> bytes: + events: Final = ( + {"type": "response.created", "response": {"id": raw_id, "status": "in_progress", "instructions": instructions}}, + {"type": "response.completed", "response": {"id": raw_id, "status": "completed", "instructions": instructions}}, + ) + return b"".join( + f"event: {event['type']}\ndata: {json.dumps(event, ensure_ascii=False)}\n\n".encode() for event in events + ) + + +@pytest.mark.asyncio +async def test_streamed_response_whose_json_holds_unicode_line_separators_is_still_owned_and_rewritten() -> None: + instructions: Final = "Answer\u2028in\u2029verse" + pc: Final = _prisma_client() + + output: Final = await _collect( + rewrite_streamed_response_ids( + stream=_chunks(_instructed_response_stream_bytes(RAW_RESPONSE_ID, instructions), 4096), + provider="openai", + method="POST", + route="/openai_passthrough/v1/responses", + user_api_key_dict=_user(), + prisma_client=pc, + ) + ) + + pc.db.litellm_managedobjecttable.upsert.assert_awaited_once() + managed_id: Final = pc.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"]["unified_object_id"] + assert output == _instructed_response_stream_bytes(managed_id, instructions) + + @pytest.mark.asyncio async def test_streamed_bytes_untouched_on_routes_without_a_response_id(): pc = _prisma_client() diff --git a/tests/unit/proxy/proxy_server/test_exception_handlers.py b/tests/unit/proxy/proxy_server/test_exception_handlers.py index f6efdecc884..b34c2d2fb08 100644 --- a/tests/unit/proxy/proxy_server/test_exception_handlers.py +++ b/tests/unit/proxy/proxy_server/test_exception_handlers.py @@ -488,22 +488,17 @@ async def test_otel_unhandled_exception_handler_reraises_http_exception_invalid( @pytest.mark.asyncio @pytest.mark.parametrize("media_type", ["application/json", "application/x-protobuf"]) @pytest.mark.parametrize("root_path", ["", "/tenant-a"]) -@pytest.mark.parametrize("native_available", [True, False]) @pytest.mark.parametrize("error", [ ProxyException("database credentials: secret", "auth_error", None, 401), HTTPException(403, "database credentials: secret"), ]) -async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native( - media_type: str, root_path: str, native_available: bool, - error: ProxyException | HTTPException, monkeypatch: pytest.MonkeyPatch, +async def test_otlp_auth_errors_hide_internal_details( + media_type: str, root_path: str, error: ProxyException | HTTPException, ) -> None: from google.rpc.status_pb2 import Status from litellm.proxy.proxy_server import otlp_http_exception_handler - from litellm.rust_bridge import loader - if not native_available: - monkeypatch.setattr(loader, "_cached_bridge", None) request: Final = Request({ "type": "http", "method": "POST", "path": root_path + "/v1/traces", "root_path": root_path, "headers": [(b"content-type", media_type.encode())], @@ -522,7 +517,7 @@ async def test_otlp_auth_errors_hide_internal_details_and_survive_missing_native else Status.FromString(response.body).message ) expected: Final = "Unauthorized" if isinstance(error, ProxyException) else "Forbidden" - assert message == (expected if native_available or media_type == "application/json" else "") + assert message == expected def _websocket(sent: list[Message]) -> WebSocket: diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 11a44a0ec16..8b1d7c1222a 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -33,9 +33,9 @@ from pydantic import BaseModel from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps +from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( ProxyStartupEvent, - initialize_shared_aiohttp_session, _resolve_pydantic_type, _resolve_typed_dict_type, cleanup_router_config_variables, @@ -43,6 +43,7 @@ from litellm.proxy.proxy_server import ( get_litellm_model_info, initialize, initialize_from_worker_config, + initialize_shared_aiohttp_session, load_from_azure_key_vault, proxy_shutdown_event, proxy_startup_event, @@ -50,8 +51,6 @@ from litellm.proxy.proxy_server import ( save_worker_config, ) -from litellm.proxy.management_helpers.audit_logs import track_audit_task - from .conftest import normalize # --------------------------------------------------------------------------- @@ -1165,7 +1164,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): @pytest.mark.asyncio async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row(): - from litellm.router_utils.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import DEFAULT_TUNING_FINGERPRINT prisma_client = MagicMock() prisma_client.db.litellm_config.find_unique = AsyncMock(return_value=None) @@ -1186,7 +1185,10 @@ async def test_tuning_baseline_v3_is_created_alongside_the_legacy_row(): @pytest.mark.asyncio async def test_scorer_baseline_upgrade_preserves_existing_routers_and_is_not_refreshed_on_restart(): - from litellm.router_utils.auto_router_tuning_baseline import mutable_tuned_identities, snapshot_tuning_baselines + from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( + mutable_tuned_identities, + snapshot_tuning_baselines, + ) deployments = [ { diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index a86314e3063..5cbf89f5e68 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/tests/unit/proxy/proxy_server/test_routes_queue.py @@ -191,6 +191,54 @@ def test_queue_chat_completions_rejects_a_request_it_cannot_queue_as_an_invalid_ assert add_request.await_count == 0 +@pytest.mark.parametrize( + ("default_priority", "priority_header", "priority_fields", "expected_priority"), + [ + (5, "0", {"priority": 5}, 0), + (5, "0", {}, 0), + (None, "high", {"priority": 1, "drop_params": True}, 1), + (None, "", {"priority": 1}, 1), + ], + ids=("header-beats-body", "header-only", "drop-params-keeps-body", "empty-header-keeps-body"), +) +def test_queue_chat_completions_queues_at_the_priority_header( + client, + auth_as, + monkeypatch: pytest.MonkeyPatch, + default_priority: int | None, + priority_header: str, + priority_fields: dict[str, object], + expected_priority: int, +) -> None: + router: Final = _queue_router(monkeypatch, default_priority=default_priority) + payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], **priority_fields} + with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request: + with auth_as(): + response: Final = client.post( + "/queue/chat/completions", json=payload, headers={"x-litellm-priority": priority_header} + ) + assert response.status_code == 200, response.text + assert [call.kwargs["request"].priority for call in add_request.await_args_list] == [expected_priority] + assert response.headers.get("x-litellm-priority") == str(expected_priority) + + +def test_queue_chat_completions_rejects_a_malformed_priority_header_like_a_malformed_body_value( + client, auth_as, monkeypatch: pytest.MonkeyPatch +) -> None: + router: Final = _queue_router(monkeypatch, default_priority=5) + payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "priority": 1} + with patch.object(router.scheduler, "add_request", wraps=router.scheduler.add_request) as add_request: + with auth_as(): + response: Final = client.post( + "/queue/chat/completions", json=payload, headers={"x-litellm-priority": "high"} + ) + assert response.status_code == 400, response.text + error: Final = response.json()["error"] + assert (error["type"], error["param"]) == ("invalid_request_error", "priority") + assert "priority must be an integer" in error["message"] + assert add_request.await_count == 0 + + def test_streamed_queue_responses_report_the_priority_the_router_queued_at( client, auth_as, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index f60c2254b41..d55ce25d878 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -363,6 +363,26 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_zerogpu_add_model_form_defaults_resolve_to_a_priced_zerogpu_model(): + import litellm + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + app_instance: Final = FastAPI() + app_instance.include_router(router) + providers: Final = TestClient(app_instance).get("/public/providers/fields").json() + + zerogpu: Final = next(p for p in providers if p["litellm_provider"] == "zerogpu") + fields_by_key: Final = {f["key"]: f for f in zerogpu["credential_fields"]} + placeholder: Final = zerogpu["default_model_placeholder"] + _, provider, _, api_base = get_llm_provider(model=placeholder) + + assert provider == "zerogpu" + assert api_base == fields_by_key["api_base"]["placeholder"] + assert litellm.model_cost[placeholder]["litellm_provider"] == "zerogpu" + assert fields_by_key["api_key"]["required"] is True + assert fields_by_key["api_key"]["field_type"] == "password" + + def test_qwen_mainland_provider_fields_carry_the_qianwen_brand(): app_instance = FastAPI() app_instance.include_router(router) diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 63e38ce533a..42acad28488 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -17,7 +17,7 @@ following the OpenAI Response API format. import asyncio import json from datetime import datetime, timezone -from typing import Any, Dict, Optional +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, Mock, patch import pytest @@ -1400,7 +1400,7 @@ def _make_sse_stream(events: list) -> Mock: async def _body_iterator(): for event in events: - yield f"data: {json.dumps(event)}" + yield f"data: {json.dumps(event, ensure_ascii=False)}" yield "data: [DONE]" mock_response = Mock() @@ -1639,6 +1639,33 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["status"] == "completed" assert final_call.kwargs["usage"] == {"input_tokens": 10, "output_tokens": 50} + @pytest.mark.asyncio + async def test_response_completed_whose_json_holds_unicode_line_separators_keeps_its_output(self) -> None: + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + item: Final = { + "id": "item_1", + "type": "message", + "content": [{"type": "output_text", "text": "Line\u2028one\u2029two"}], + } + events: Final = [ + {"type": "response.completed", "response": {"id": "resp_123", "status": "completed", "output": [item]}} + ] + handler: Final = AsyncMock(spec=ResponsePollingHandler) + + with patch( + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock(return_value=_make_sse_stream(events)) + await background_streaming_task(**_make_background_streaming_kwargs("poll_u2028", handler)) + + calls: Final = handler.update_state.call_args_list + outputs: Final = [call.kwargs["output"] for call in calls if "output" in call.kwargs] + assert outputs[-1] == [item] + assert calls[-1].kwargs["status"] == "completed" + @pytest.mark.asyncio async def test_fallback_status_derived_from_event_type_when_status_field_missing( self, diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 7e8db33095f..9e149bde2c8 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -35,6 +35,7 @@ from litellm.constants import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ProxyConfig, @@ -574,6 +575,168 @@ class TestProxyBaseLLMRequestProcessing: assert returned_data["cache"] == expected_cache_controls assert returned_data["fallbacks"] == [{"model": "gpt-4o", "cache": expected_fallback_cache_controls}] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("path", "route_type", "has_router", "priority_header", "body_priority", "expected_priority"), + [ + ("/v1/chat/completions", "acompletion", True, "0", 5, 0), + ("/v1/completions", "atext_completion", True, "0", 5, 0), + ("/v1/chat/completions", "acompletion", True, "1", None, 1), + ("/v1/completions", "atext_completion", True, "1", None, 1), + ("/v1/chat/completions", "acompletion", True, None, 5, 5), + ("/v1/chat/completions", "acompletion", True, "", 5, 5), + ("/cursor/chat/completions", "aresponses", True, "0", None, None), + ("/v1/responses", "aresponses", True, "0", 5, 5), + ("/v1/embeddings", "aembedding", True, "0", None, None), + ("/v1/chat/completions", "acompletion", False, "0", None, None), + ("/v1/chat/completions", "acompletion", False, "0", 5, 5), + ], + ids=( + "chat-header-beats-body", + "text-completion-header-beats-body", + "chat-header-only", + "text-completion-header-only", + "chat-body-only", + "chat-empty-header-keeps-body", + "cursor-responses-route-ignores-header", + "responses-keeps-body", + "embeddings-ignores-header", + "routerless-ignores-header", + "routerless-keeps-body", + ), + ) + async def test_common_processing_pre_call_logic_applies_priority_header_only_where_the_router_schedules( + self, + path: str, + route_type: Literal["acompletion", "atext_completion", "aresponses", "aembedding"], + has_router: bool, + priority_header: str | None, + body_priority: int | None, + expected_priority: str | int | None, + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + header_entries: Final[tuple[tuple[bytes, bytes], ...]] = ( + () if priority_header is None else ((b"x-litellm-priority", priority_header.encode()),) + ) + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": path, + "scheme": "http", + "server": ("localhost", 4000), + "client": ("127.0.0.1", 12345), + "query_string": b"", + "headers": ((b"content-type", b"application/json"), *header_entries), + } + ) + router: Final = ( + Router( + model_list=[ + { + "model_name": "qa-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test", "mock_response": "hi"}, + } + ] + ) + if has_router + else None + ) + processing_obj: Final = ProxyBaseLLMRequestProcessing( + data={ + "model": "qa-model", + "messages": [{"role": "user", "content": "hi"}], + **({"priority": body_priority} if body_priority is not None else {}), + } + ) + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_config=RealProxyConfig(), + route_type=route_type, + llm_router=router, + ) + + assert ("priority" in returned_data, returned_data.get("priority")) == ( + expected_priority is not None, + expected_priority, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("requested_model", "deployment_drops_params", "body_priority", "expected_priority"), + [ + ("qa-model", True, 5, 5), + ("qa-model", True, None, None), + ("key-alias", True, 5, 5), + ("qa-model", False, 5, "high"), + ], + ids=( + "drop-params-keeps-body", + "drop-params-without-body-leaves-default", + "drop-params-read-from-the-aliased-deployment", + "no-drop-params-passes-header-on-to-400", + ), + ) + async def test_common_processing_pre_call_logic_drops_only_a_malformed_priority_header_under_drop_params( + self, + requested_model: str, + deployment_drops_params: bool, + body_priority: int | None, + expected_priority: str | int | None, + ) -> None: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "scheme": "http", + "server": ("localhost", 4000), + "client": ("127.0.0.1", 12345), + "query_string": b"", + "headers": ((b"content-type", b"application/json"), (b"x-litellm-priority", b"high")), + } + ) + router: Final = Router( + model_list=[ + { + "model_name": "qa-model", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-test", + "mock_response": "hi", + "drop_params": deployment_drops_params, + }, + } + ] + ) + processing_obj: Final = ProxyBaseLLMRequestProcessing( + data={ + "model": requested_model, + "messages": [{"role": "user", "content": "hi"}], + **({"priority": body_priority} if body_priority is not None else {}), + } + ) + returned_data, _ = await processing_obj.common_processing_pre_call_logic( + request=request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", aliases={"key-alias": "qa-model"}), + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + proxy_config=RealProxyConfig(), + route_type="acompletion", + llm_router=router, + ) + + assert ("priority" in returned_data, returned_data.get("priority")) == ( + expected_priority is not None, + expected_priority, + ) + @pytest.mark.asyncio @pytest.mark.parametrize("safe_memory_mode", [False, True]) @pytest.mark.parametrize( @@ -10887,12 +11050,19 @@ class TestResolveLitellmCallId: class _FailureHookRecorder: """Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it.""" - def __init__(self, raises: Optional[Exception] = None): + def __init__(self, raises: Optional[Exception] = None, gate: Optional[asyncio.Event] = None): self.calls = [] + self.entered = asyncio.Event() + self.reported = asyncio.Event() self._raises = raises + self._gate = gate async def post_call_failure_hook(self, **kwargs): + self.entered.set() + if self._gate is not None: + await self._gate.wait() self.calls.append(kwargs) + self.reported.set() if self._raises is not None: raise self._raises @@ -10993,6 +11163,171 @@ class TestDetachedStreamFailureHook: assert [call["original_exception"] for call in recorder.calls] == [failure] + @staticmethod + def _armed_stream(recorder: _FailureHookRecorder, iterator_hook_for, response: Optional[MagicMock] = None): + from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth + + logging_obj = LiteLLMLoggingObj( + model="gpt-6.1-sol", + messages=[{"role": "user", "content": "Write a story."}], + stream=True, + call_type="anthropic_messages", + start_time=None, + litellm_call_id="call-lit6871", + function_id="function-lit6871", + ) + user_api_key_dict = RealUserAPIKeyAuth(api_key="sk-test") + processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj}) + ProxyLogging._callback_capabilities_cache.clear() + proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) + proxy_logging_obj.post_call_failure_hook = recorder.post_call_failure_hook + proxy_logging_obj.async_post_call_streaming_iterator_hook = iterator_hook_for(logging_obj) + processing_obj._arm_detached_stream_failure_hook( + logging_obj=logging_obj, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + stream = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( + response=MagicMock() if response is None else response, + user_api_key_dict=user_api_key_dict, + request_data=processing_obj.data, + proxy_logging_obj=proxy_logging_obj, + serialize_chunk=lambda chunk: "data: chunk\n\n", + serialize_error=lambda error: "data: error\n\n", + ) + return logging_obj, stream + + @staticmethod + def _hook_reporting(upstream_failure: Exception, after_report=None): + def iterator_hook_for(logging_obj): + async def iterator_hook(**_kwargs): + yield b"event: message_start\n\n" + await logging_obj.on_detached_stream_failure(upstream_failure) + yield b"event: error\n\n" + yield b"event: message_stop\n\n" + if after_report is not None: + raise after_report + + return iterator_hook + + return iterator_hook_for + + @pytest.mark.asyncio + async def test_iterator_hook_failure_after_a_bridged_failure_is_the_one_the_failure_hook_gets(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed") + hook_failure = RuntimeError("iterator hook rejected the stream") + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure)) + + frames = [frame async for frame in stream] + + assert frames == ["data: chunk\n\n"] * 3 + ["data: error\n\n"] + assert [call["original_exception"] for call in recorder.calls] == [hook_failure] + + @pytest.mark.asyncio + async def test_failure_hook_replacing_the_iterator_hook_failure_aborts_the_stream(self): + replacement = HTTPException(status_code=418, detail="rejected by a custom failure hook") + recorder = _FailureHookRecorder(raises=replacement) + upstream_failure = RuntimeError("upstream Responses stream failed") + hook_failure = RuntimeError("iterator hook rejected the stream") + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure, after_report=hook_failure)) + + frames = [await stream.__anext__() for _ in range(3)] + with pytest.raises(HTTPException) as raised: + await stream.__anext__() + + assert raised.value is replacement + assert frames == ["data: chunk\n\n"] * 3 + assert [call["original_exception"] for call in recorder.calls] == [hook_failure] + + @pytest.mark.asyncio + async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_after_the_stream_ends(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed") + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure)) + + frames = [] + async for frame in stream: + frames.append(frame) + assert recorder.calls == [] + + assert frames == ["data: chunk\n\n"] * 3 + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + + @pytest.mark.asyncio + async def test_failure_report_and_stream_cleanup_survive_the_client_task_being_cancelled_mid_report(self): + gate = asyncio.Event() + recorder = _FailureHookRecorder(gate=gate) + upstream_failure = RuntimeError("upstream Responses stream failed") + response = MagicMock() + response.aclose = AsyncMock() + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure), response=response) + + async def client_reads_the_whole_stream(): + return [frame async for frame in stream] + + client = asyncio.create_task(client_reads_the_whole_stream()) + await asyncio.wait_for(recorder.entered.wait(), timeout=5) + client.cancel() + with pytest.raises(asyncio.CancelledError): + await client + + response.aclose.assert_awaited_once() + gate.set() + await asyncio.wait_for(recorder.reported.wait(), timeout=5) + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + + @pytest.mark.asyncio + async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_when_the_client_leaves_early(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed") + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure)) + + await stream.__anext__() + await stream.__anext__() + await stream.aclose() + + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + + @pytest.mark.asyncio + async def test_failure_reported_after_the_stream_ends_still_reaches_the_failure_hook(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed after the client stream closed") + + async def iterator_hook(**_kwargs): + yield b"event: message_start\n\n" + + logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook) + + frames = [frame async for frame in stream] + await logging_obj.on_detached_stream_failure(upstream_failure) + + assert frames == ["data: chunk\n\n"] + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + + @pytest.mark.asyncio + async def test_failure_reported_by_an_iterator_hook_is_not_reported_again_when_the_stream_closes(self): + recorder = _FailureHookRecorder() + upstream_failure = RuntimeError("upstream Responses stream failed") + hook_failure = RuntimeError("iterator hook rejected the error frame") + + async def iterator_hook(**_kwargs): + yield b"event: message_start\n\n" + raise hook_failure + + logging_obj, stream = self._armed_stream(recorder, lambda _logging_obj: iterator_hook) + + frames = [frame async for frame in stream] + await logging_obj.on_detached_stream_failure(upstream_failure) + + assert frames == ["data: chunk\n\n", "data: error\n\n"] + assert [call["original_exception"] for call in recorder.calls] == [hook_failure] + class TestPostCallMaskedOutputReachesDeferredLogging: @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 6044475652c..9e31d3774f9 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1494,7 +1494,7 @@ async def test_route_request_keeps_key_fallbacks_out_of_pass_through_provider_bo ("avideo_create_character", "name", "/videos/characters"), ("acreate_container", "name", "/containers"), ("aupload_container_file", "file", "/containers/{container_id}/files"), - ("acreate_agent", "name", "/v1beta/agents"), + ("acreate_agent", "id", "/v1beta/agents"), ("acreate_eval", "data_source_config", "/evals"), ("acreate_run", "data_source", "/evals/{eval_id}/runs"), ], @@ -1699,7 +1699,7 @@ def test_required_present_body_param_with_none_router_wide_default_still_raises( ("asearch", {"model": "search-model"}, "query"), ("acreate_eval", {"model": "eval-model"}, "data_source_config"), ("acreate_run", {"model": "eval-model"}, "data_source"), - ("acreate_agent", {"model": "agent-model"}, "name"), + ("acreate_agent", {"model": "agent-model"}, "id"), ("avector_store_search", {}, "query"), ], ) @@ -1842,6 +1842,8 @@ def test_raise_if_required_body_param_missing_names_each_missing_param(route_typ ("acreate_eval", {"data_source_config": {"type": "custom"}, "testing_criteria": []}), ("acreate_interaction", {"model": "gemini-3-pro-preview", "input": "hi"}), ("acreate_interaction", {"agent": "deep-research", "input": "hi"}), + ("acreate_agent", {"id": "sentinel"}), + ("acreate_agent", {"name": "sentinel"}), ("aimage_generation", {"model": "gpt-image-1", "prompt": "a cat"}), ("aspeech", {"model": "gpt-4o-mini-tts", "input": "hi", "voice": "alloy"}), ("amoderation", {"model": "omni-moderation-latest", "input": ""}), diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index 62ca71b6cf6..7cff44da985 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -30,7 +30,6 @@ from litellm.proxy.auth.authorization import OwnedRows, ReadScope from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.tracing_runtime import manage_tracing, provide_storage -from litellm.rust_bridge import loader from litellm.tracing.errors import TraceChanged from litellm.tracing.generated.models import TraceQueryHelp from litellm.tracing.generated.responses import TraceSQLResponse diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index ed90acbd2fa..0b971bc1696 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1288,6 +1288,106 @@ async def test_reasoning_item_closes_before_a_tool_call_from_the_same_chunk_open ] +def _text_opening_a_tool_call_chunk(content: str, arguments: str = "") -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role="assistant", + content=content, + tool_calls=[ + { + "id": "call_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": arguments}, + "index": 0, + } + ], + ), + finish_reason=None, + ) + ], + ) + + +def _tool_arguments_chunk(arguments: str, finish_reason: str | None = None) -> ModelResponseStream: + return ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-haiku-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + tool_calls=[{"type": "function", "function": {"arguments": arguments}, "index": 0}], + ), + finish_reason=finish_reason, + ) + ], + ) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_call_opened_in_a_text_chunk_streams_each_argument_delta(sync_mode: bool) -> None: + argument_pieces: Final = ('{"', "city", '"', ': "', "Paris", '"', "}") + iterator: Final = _build_iterator( + [ + _text_opening_a_tool_call_chunk(""), + *(_tool_arguments_chunk(piece) for piece in argument_pieces), + _chunk("", finish_reason="tool_calls"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "" + assert [ + event.delta + for event in events + if getattr(event, "type", None) == ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA + ] == list(argument_pieces) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_tool_call_sharing_a_chunk_with_text_opens_after_the_whole_message(sync_mode: bool) -> None: + iterator: Final = _build_iterator( + [ + _text_opening_a_tool_call_chunk("Checking ", arguments='{"city": "Paris"}'), + _chunk("the "), + _chunk("weather now."), + _chunk("", finish_reason="tool_calls"), + ] + ) + + events: Final = await _collect_events(iterator, sync_mode) + + event_types: Final = [getattr(event, "type", None) for event in events] + function_call_added_at: Final = next( + position + for position, event in enumerate(events) + if event_types[position] == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED and event.item.type == "function_call" + ) + streamed: Final = [ + event + for event in events + if getattr(event, "type", None) + in (ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA) + ] + assert ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA not in event_types[function_call_added_at:] + assert [event.sequence_number for event in streamed] == sorted(event.sequence_number for event in streamed) + assert _event_text(events, ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA) == "Checking the weather now." + assert _event_text(events, ResponsesAPIStreamEvents.FUNCTION_CALL_ARGUMENTS_DELTA) == '{"city": "Paris"}' + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_tool_only_stream_puts_the_function_call_at_output_index_zero(sync_mode: bool): diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 4cbcc6c1519..ec0be896095 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -2668,6 +2668,65 @@ async def test_stream_auto_execute_bills_the_paid_first_round_on_its_success_row assert _request_failure_row(logging_obj, recorder.logged[0]).get("response_cost", 0.0) == 0.0 +class _CutAfterStreamedText(httpx.AsyncByteStream): + async def __aiter__(self): + opened: Final = {**json.loads(json.dumps(_ROUND_TWO_BODY, default=dict)), "status": "in_progress"} + yield _sse( + ( + {"type": "response.created", "response": opened, "sequence_number": 0}, + { + "type": "response.output_text.delta", + "item_id": "msg_round_two", + "output_index": 0, + "content_index": 0, + "delta": "PING " * 40, + "sequence_number": 1, + }, + ) + ) + raise httpx.ReadError("connection reset mid-stream") + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_a_follow_up_cut_after_streamed_text_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), httpx.Response(200, stream=_CutAfterStreamedText(), headers=SSE_HEADERS)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-cut-after-text") + + with pytest.raises(httpx.ReadError): + await _stream_auto_execute(logging_obj) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert len(recorder.failed) == 1 + failed: Final = recorder.failed[0] + assert failed["completion_tokens"] > 0 + priced_with_its_own_margin: Final = litellm.completion_cost( + completion_response=ModelResponse( + model="gpt-5.6", + usage=Usage( + prompt_tokens=failed["prompt_tokens"], + completion_tokens=failed["completion_tokens"], + total_tokens=failed["prompt_tokens"] + failed["completion_tokens"], + ), + ), + model="gpt-5.6", + custom_llm_provider="openai", + ) + assert failed["response_cost"] == pytest.approx(priced_with_its_own_margin - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(failed) + + @pytest.mark.asyncio async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index 37a58c32f38..ae688ff9173 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -9,20 +9,23 @@ crashed with exit 139 whenever any CustomLogger was registered). import asyncio import json import time -from collections.abc import Mapping -from typing import Final +from collections.abc import AsyncIterator, Mapping +from datetime import datetime +from typing import Final, cast import httpx import pytest import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import thread_pool_executor as thread_pool_executor_module -from litellm.responses import streaming_iterator as responses_streaming_iterator_module from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.responses import streaming_iterator as responses_streaming_iterator_module from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse +from litellm.types.utils import StandardLoggingPayload class RecordingCustomLogger(CustomLogger): @@ -169,6 +172,240 @@ async def test_sync_callbacks_run_only_after_async_handler_completes(recording_e assert submit_times[0] >= recorder.async_hook_finished +class FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failure_payloads: tuple[StandardLoggingPayload, ...] = () + self.success_payloads: tuple[StandardLoggingPayload, ...] = () + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failure_payloads = ( + *self.failure_payloads, + cast(StandardLoggingPayload, kwargs["standard_logging_object"]), + ) + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.success_payloads = ( + *self.success_payloads, + cast(StandardLoggingPayload, kwargs["standard_logging_object"]), + ) + + +class _ChunkRecorder: + def __init__(self) -> None: + self.chunks: tuple[ResponsesAPIStreamingResponse, ...] = () + + async def drain(self, iterator: ResponsesAPIStreamingIterator) -> None: + async for chunk in iterator: + self.chunks = (*self.chunks, chunk) + + +_SEA_DELTAS = ("The sea ", "is wide ", "and deep.") +_SEA_PROMPT = "Write a 300 word story about the sea." +_UNREACHABLE_IMAGE_INPUT: list[dict[str, object]] = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Describe this picture."}, + {"type": "input_image", "image_url": "http://images.invalid/photo.png", "detail": "high"}, + ], + } +] + + +_TOOL_ARGUMENT_DELTAS = ('{"city": ', '"Lisbon", ', '"unit": "celsius"}') + + +class _UpstreamThatTimesOutAfterThreeDeltas: + headers: dict[str, str] = {} + + def __init__(self, event_type: str = "response.output_text.delta", deltas: tuple[str, ...] = _SEA_DELTAS) -> None: + self.event_type = event_type + self.deltas = deltas + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + for sequence_number, delta in enumerate(self.deltas): + event = { + "type": self.event_type, + "item_id": "item_1", + "output_index": 0, + "content_index": 0, + "delta": delta, + "sequence_number": sequence_number, + } + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +class _UpstreamThatFailsWithUsageAfterThreeDeltas: + headers: dict[str, str] = {} + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + for sequence_number, delta in enumerate(_SEA_DELTAS): + event = { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": delta, + "sequence_number": sequence_number, + } + yield f"data: {json.dumps(event)}\n\n".encode() + failed = { + "type": "response.failed", + "sequence_number": len(_SEA_DELTAS), + "response": { + "id": "resp_1", + "object": "response", + "created_at": 1700000000, + "status": "failed", + "model": "gpt-5.4-nano", + "output": [], + "error": {"code": "server_error", "message": "The model had an internal error"}, + "usage": {"input_tokens": 400, "output_tokens": 9, "total_tokens": 409}, + }, + } + yield f"data: {json.dumps(failed)}\n\n".encode() + + +async def _failure_payload_after_mid_stream_timeout( + monkeypatch: pytest.MonkeyPatch, + request_input: str | list[dict[str, object]], + upstream: _UpstreamThatTimesOutAfterThreeDeltas | None = None, + background_fields: Mapping[str, object] | None = None, +) -> tuple[FailureRecorder, StandardLoggingPayload]: + upstream_response: Final = upstream or _UpstreamThatTimesOutAfterThreeDeltas() + recorder = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + logging_obj = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator = ResponsesAPIStreamingIterator( + response=cast(httpx.Response, upstream_response), + model="gpt-5.4-nano", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + custom_llm_provider="openai", + request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, **(background_fields or {})}, + ) + received = _ChunkRecorder() + + with pytest.raises(litellm.Timeout, match="Timeout on reading data from socket"): + await received.drain(iterator) + assert [getattr(chunk, "delta", None) for chunk in received.chunks] == list(upstream_response.deltas) + + await iterator._await_pending_logging() + (payload,) = recorder.failure_payloads + return recorder, payload + + +@pytest.mark.asyncio +async def test_mid_stream_read_timeout_logs_failure_with_partial_usage_and_no_success( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT) + + assert payload["status"] == "failure" + assert payload["prompt_tokens"] == litellm.token_counter( + model="gpt-5.4-nano", messages=[{"role": "user", "content": _SEA_PROMPT}] + ) + assert payload["completion_tokens"] == litellm.token_counter( + model="gpt-5.4-nano", text="".join(_SEA_DELTAS), count_response_tokens=True + ) + assert payload["response_cost"] > 0 + assert recorder.success_payloads == () + + +@pytest.mark.asyncio +async def test_mid_stream_read_timeout_bills_the_tool_call_arguments_already_streamed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + upstream: Final = _UpstreamThatTimesOutAfterThreeDeltas( + event_type="response.function_call_arguments.delta", deltas=_TOOL_ARGUMENT_DELTAS + ) + _, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT, upstream) + + assert payload["status"] == "failure" + assert payload["completion_tokens"] == litellm.token_counter( + model="gpt-5.4-nano", text="".join(_TOOL_ARGUMENT_DELTAS), count_response_tokens=True + ) + assert payload["response_cost"] > 0 + + +@pytest.mark.asyncio +async def test_mid_stream_failure_of_a_background_response_bills_nothing_since_retrieving_it_bills_the_whole_run( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, payload = await _failure_payload_after_mid_stream_timeout( + monkeypatch, _SEA_PROMPT, background_fields={"background": True} + ) + + assert payload["status"] == "failure" + assert payload["completion_tokens"] == 0 + assert payload["response_cost"] == 0 + + +@pytest.mark.asyncio +async def test_mid_stream_failure_usage_estimate_never_fetches_image_dimensions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _UNREACHABLE_IMAGE_INPUT) + + assert payload["status"] == "failure" + assert payload["prompt_tokens"] == litellm.token_counter( + model="gpt-5.4-nano", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this picture."}, + {"type": "image_url", "image_url": {"url": "http://images.invalid/photo.png", "detail": "high"}}, + ], + } + ], + use_default_image_token_count=True, + ) + assert payload["completion_tokens"] > 0 + + +@pytest.mark.asyncio +async def test_mid_stream_failure_keeps_the_usage_the_provider_reported(monkeypatch: pytest.MonkeyPatch) -> None: + recorder = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + logging_obj = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator = ResponsesAPIStreamingIterator( + response=cast(httpx.Response, _UpstreamThatFailsWithUsageAfterThreeDeltas()), + model="gpt-5.4-nano", + responses_api_provider_config=OpenAIResponsesAPIConfig(), + logging_obj=logging_obj, + custom_llm_provider="openai", + request_data={"model": "gpt-5.4-nano", "input": _SEA_PROMPT, "stream": True}, + ) + + with pytest.raises(MidStreamFallbackError, match="The model had an internal error"): + await _ChunkRecorder().drain(iterator) + await iterator._await_pending_logging() + + (payload,) = recorder.failure_payloads + assert (payload["prompt_tokens"], payload["completion_tokens"]) == (400, 9) + + _CREATED_RESPONSE_WITH_PRIORITY_TIER: Final[Mapping[str, object]] = { "id": "resp_partial", "object": "response", @@ -211,17 +448,135 @@ def _completed_after(text: str) -> Mapping[str, object]: } -def _live_iterator(logging_obj: LitellmLogging) -> ResponsesAPIStreamingIterator: +def _live_iterator( + logging_obj: LitellmLogging, upstream: object = None, background: bool = False +) -> ResponsesAPIStreamingIterator: return ResponsesAPIStreamingIterator( - response=httpx.Response(200), + response=cast(httpx.Response, upstream or httpx.Response(200)), model="gpt-5.4-nano", responses_api_provider_config=OpenAIResponsesAPIConfig(), logging_obj=logging_obj, - request_data={"model": "gpt-5.4-nano", "input": "tell me a story"}, + request_data={"model": "gpt-5.4-nano", "input": "tell me a story", "background": background}, custom_llm_provider="openai", ) +class _UpstreamThatTimesOutAfterRelayingText: + headers: dict[str, str] = {} + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)): + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +class _UpstreamThatCompletesThenTimesOut: + headers: dict[str, str] = {} + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + events: Final = ( + _CREATED_WITH_PRIORITY_TIER, + *(_text_delta(delta, sequence_number) for sequence_number, delta in enumerate(_SEA_DELTAS, start=1)), + _completed_after("".join(_SEA_DELTAS)), + ) + for event in events: + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +@pytest.mark.asyncio +async def test_timeout_after_completed_response_ends_the_stream_and_bills_it_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + recorder: Final = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + logging_obj: Final = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatCompletesThenTimesOut()) + + chunks: Final = _ChunkRecorder() + await chunks.drain(iterator) + await iterator._await_pending_logging() + + assert chunks.chunks[-1].type == "response.completed" + assert len(recorder.success_payloads) == 1 + success_payload: Final = recorder.success_payloads[0] + response_payload: Final = success_payload["response"] + assert isinstance(response_payload, dict) + usage: Final = cast(dict[str, object], response_payload["usage"]) + assert (usage["prompt_tokens"], usage["completion_tokens"], usage["total_tokens"]) == (5, 7, 12) + assert recorder.failure_payloads == () + assert logging_obj.model_call_details.get("combined_usage_object") is None + + +class _UpstreamThatTimesOutAfterCreated: + headers: dict[str, str] = {} + + def __init__(self, **created_fields: object) -> None: + self.created_fields = created_fields + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + created: Final = { + **_CREATED_WITH_PRIORITY_TIER, + "response": {**_CREATED_RESPONSE_WITH_PRIORITY_TIER, **self.created_fields}, + } + for event in (created, _text_delta("Once upon a time", 1)): + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +async def _mid_stream_failure_payload( + monkeypatch: pytest.MonkeyPatch, requested_tier: str | None, **created_fields: object +) -> StandardLoggingPayload: + recorder: Final = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + logging_obj: Final = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True} + if requested_tier is None + else {"stream": True, "service_tier": requested_tier}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamThatTimesOutAfterCreated(**created_fields)) + + with pytest.raises(litellm.Timeout): + await _ChunkRecorder().drain(iterator) + await iterator._await_pending_logging() + + (payload,) = recorder.failure_payloads + return payload + + +@pytest.mark.asyncio +async def test_mid_stream_failure_bills_the_relayed_text_at_the_tier_the_provider_served( + monkeypatch: pytest.MonkeyPatch, +) -> None: + flex_requested: Final = await _mid_stream_failure_payload(monkeypatch, "flex", service_tier="default") + nothing_requested: Final = await _mid_stream_failure_payload(monkeypatch, None, service_tier="default") + + assert nothing_requested["response_cost"] > 0 + assert flex_requested["response_cost"] == nothing_requested["response_cost"] + + +@pytest.mark.asyncio +async def test_mid_stream_failure_bills_nothing_when_the_provider_runs_the_response_in_background( + monkeypatch: pytest.MonkeyPatch, +) -> None: + payload: Final = await _mid_stream_failure_payload(monkeypatch, None, background=True) + + assert payload["status"] == "failure" + assert (payload["completion_tokens"], payload["response_cost"]) == (0, 0) + + def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier(): iterator: Final = _live_iterator(_make_logging_obj()) for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)): @@ -263,3 +618,25 @@ async def test_partial_response_on_disconnect_is_none_when_there_is_nothing_to_b assert iterator.partial_response_on_disconnect() is None await iterator._await_pending_logging() + + +@pytest.mark.asyncio +async def test_partial_response_on_disconnect_is_none_once_the_failure_log_billed_the_relayed_text(): + iterator: Final = _live_iterator(_make_logging_obj(), upstream=_UpstreamThatTimesOutAfterRelayingText()) + + with pytest.raises(litellm.Timeout): + await _ChunkRecorder().drain(iterator) + + assert iterator.partial_response_on_disconnect() is None + await iterator._await_pending_logging() + + +def test_partial_response_on_disconnect_reports_the_relayed_text_of_a_background_response(): + iterator: Final = _live_iterator(_make_logging_obj(), background=True) + for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)): + iterator._process_chunk(json.dumps(event)) + + incomplete_event: Final = iterator.partial_response_on_disconnect() + + assert incomplete_event is not None + assert incomplete_event.response.output[0].content[0].text == "Once upon a time" diff --git a/tests/unit/responses/test_rust_bridge_websocket.py b/tests/unit/responses/test_rust_bridge_websocket.py index 00ae5eb970f..57b161d2147 100644 --- a/tests/unit/responses/test_rust_bridge_websocket.py +++ b/tests/unit/responses/test_rust_bridge_websocket.py @@ -55,21 +55,6 @@ async def test_adapter_raises_clean_close_when_rust_connection_ends() -> None: await adapter.recv() -@pytest.mark.asyncio -async def test_bridge_unavailable_returns_none(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(responses_websocket, "_STATE", responses_websocket._RustResponsesWebSocketState()) - monkeypatch.setattr(responses_websocket, "get_native_bridge", lambda: None) - - assert ( - await responses_websocket.connect( - url="wss://example.test/responses", - headers={}, - timeout=None, - ) - is None - ) - - @pytest.mark.asyncio async def test_enabled_bridge_connects_and_adapts_socket( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 6c3e6177503..664f4dba163 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -210,9 +210,11 @@ _TRAILING_ERRORS = [ httpx.RemoteProtocolError("peer closed connection without sending complete message body"), ] +_TRAILING_ERRORS_AFTER_COMPLETED = [*_TRAILING_ERRORS, httpx.ReadTimeout("Timeout on reading data from socket")] + @pytest.mark.asyncio -@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS_AFTER_COMPLETED, ids=type) async def test_transport_error_after_completed_event_ends_stream_cleanly(trailing_error): """A sloppy connection close after `response.completed` must not turn a complete stream into an error (regression guard for the transport no longer @@ -243,7 +245,7 @@ async def test_transport_error_before_completed_event_raises(): pass -@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) +@pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS_AFTER_COMPLETED, ids=type) def test_sync_transport_error_after_completed_event_ends_stream_cleanly(trailing_error): iterator = _make_sync_iterator( sse_events=_COMPLETE_STREAM_EVENTS, @@ -1205,7 +1207,9 @@ async def test_completed_event_survives_a_failing_usage_estimate(): falls back to usage None.""" malformed_input: Final = [{"type": "message", "role": "user", "content": 42}] with pytest.raises(ValueError, match="Invalid content type"): - _estimate_usage_from_text("gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world") + _estimate_usage_from_text( + "gpt-4o-mini", malformed_input, {"input": malformed_input}, "hello world" + ) response = _responses_api_response_without_usage() iterator = _make_iterator( diff --git a/tests/unit/router_utils/test_auto_router_model_naming.py b/tests/unit/router_strategy/test_auto_router_model_naming.py similarity index 99% rename from tests/unit/router_utils/test_auto_router_model_naming.py rename to tests/unit/router_strategy/test_auto_router_model_naming.py index 1cfa3aa5c91..44fcfdaf5be 100644 --- a/tests/unit/router_utils/test_auto_router_model_naming.py +++ b/tests/unit/router_strategy/test_auto_router_model_naming.py @@ -3,9 +3,7 @@ from typing import Final import pytest -from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets -from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS -from litellm.router_utils.auto_router_model_naming import ( +from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, capability_limit_violation, carries_complexity_router_settings, @@ -18,6 +16,8 @@ from litellm.router_utils.auto_router_model_naming import ( validate_complexity_router_config_write, validate_strategy_router_model_write, ) +from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets +from litellm.router_strategy.complexity_router.jev_classifier import DEFAULT_JEV_INSTRUCTIONS COMPLEXITY_FIELDS = frozenset({"complexity_router_config"}) SEMANTIC_FIELDS = frozenset({"auto_router_config", "auto_router_default_model", "auto_router_embedding_model"}) @@ -596,7 +596,7 @@ def test_custom_classifier_prompt_capability(config: Mapping[str, object], expec ], ) def test_is_complexity_router_model(model: str | None, expected: bool) -> None: - from litellm.router_utils.auto_router_model_naming import is_complexity_router_model + from litellm.router_strategy.auto_router.auto_router_model_naming import is_complexity_router_model assert is_complexity_router_model(model) is expected diff --git a/tests/unit/router_utils/test_auto_router_tuning_baseline.py b/tests/unit/router_strategy/test_auto_router_tuning_baseline.py similarity index 99% rename from tests/unit/router_utils/test_auto_router_tuning_baseline.py rename to tests/unit/router_strategy/test_auto_router_tuning_baseline.py index d0d98ce7042..4ee219a7784 100644 --- a/tests/unit/router_utils/test_auto_router_tuning_baseline.py +++ b/tests/unit/router_strategy/test_auto_router_tuning_baseline.py @@ -7,8 +7,7 @@ from typing import Final import pytest -from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig -from litellm.router_utils.auto_router_tuning_baseline import ( +from litellm.router_strategy.auto_router.auto_router_tuning_baseline import ( DEFAULT_TUNING_FINGERPRINT, HEURISTIC_V1_TUNING_FIELDS, heuristic_v1_router_fingerprint, @@ -19,6 +18,7 @@ from litellm.router_utils.auto_router_tuning_baseline import ( tuning_limit_violation, tuning_quota_violation, ) +from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig _TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"} _ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"} diff --git a/tests/unit/router_utils/test_baseline_request.py b/tests/unit/router_strategy/test_baseline_request.py similarity index 95% rename from tests/unit/router_utils/test_baseline_request.py rename to tests/unit/router_strategy/test_baseline_request.py index 37d21597752..eb785a41d63 100644 --- a/tests/unit/router_utils/test_baseline_request.py +++ b/tests/unit/router_strategy/test_baseline_request.py @@ -1,6 +1,6 @@ from typing import Final -from litellm.router_utils.baseline_request import baseline_request, capture_baseline_parameters +from litellm.router_strategy.auto_router.baseline_request import baseline_request, capture_baseline_parameters def test_baseline_snapshot_owns_nested_caller_settings_and_overrides_routed_settings() -> None: diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index ca39c7f5c8f..018f93779a6 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -24,14 +24,6 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm import Router -from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import StandardLoggingRoutingDecision -from litellm.router_utils.auto_router_model_naming import ( - CUSTOMIZATION_CAPABILITY, - GATED_AUTO_ROUTER_CAPABILITIES, - HEURISTIC_V2_CAPABILITY, - count_capability_routers, -) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache @@ -42,7 +34,19 @@ from litellm.constants import ( RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY, ) +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.router import as_output_cap, stamped_client_ceiling +from litellm.router_strategy.auto_router.auto_router_model_naming import ( + CUSTOMIZATION_CAPABILITY, + GATED_AUTO_ROUTER_CAPABILITIES, + HEURISTIC_V2_CAPABILITY, + count_capability_routers, +) +from litellm.router_strategy.complexity_router.capability_classifier import ( + CAPABILITY_CLASSIFIER_SYSTEM_PROMPT, + CapabilityClassifierVerdict, +) from litellm.router_strategy.complexity_router.complexity_router import ( _CLASSIFICATION_CURRENT_MESSAGE_ONLY, _CLASSIFICATION_WITH_CONVERSATION, @@ -58,18 +62,14 @@ from litellm.router_strategy.complexity_router.complexity_router import ( classification_system_prompt, custom_tier_classification_prompt, ) -from litellm.router_strategy.complexity_router.capability_classifier import ( - CAPABILITY_CLASSIFIER_SYSTEM_PROMPT, - CapabilityClassifierVerdict, -) from litellm.router_strategy.complexity_router.config import ( - CapabilityCalibrationConfig, - CapabilityClassifierConfig, DEFAULT_CLASSIFICATION_RUBRIC, DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, DEFAULT_COMPLEXITY_CONFIG, DEFAULT_TECHNICAL_KEYWORDS, TIER_SEVERITY_ORDER, + CapabilityCalibrationConfig, + CapabilityClassifierConfig, ClassificationRubric, ClassifierLLMConfig, ComplexityRouterConfig, @@ -83,10 +83,12 @@ from litellm.router_strategy.complexity_router.jev_classifier import ( JevUsage, ) from litellm.router_strategy.complexity_router.llm_v2 import LLM_V2_PROMPT_VERSION +from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig from litellm.router_strategy.complexity_router.tier_predictor import ( TierGlobalStatistic, TrainedTierArtifact, ) +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import ( Deployment, LiteLLM_Params, @@ -94,10 +96,7 @@ from litellm.types.router import ( RouterErrors, TaggedPreRoutingStrategy, ) -from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.router_strategy.complexity_router.request_models import RequestComplexityRouterConfig -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler - +from litellm.types.utils import StandardLoggingRoutingDecision requires_semantic_router = pytest.mark.skipif( sys.version_info >= (3, 14), reason="The semantic-router extra excludes Python 3.14" @@ -16299,7 +16298,6 @@ class TestTierHealthFailover: tests can prove the probe passes a kwargs copy and forwards the prompt arguments. """ import litellm as litellm_module - from litellm.types.router import RouterRateLimitError probed_kwargs = [] diff --git a/tests/unit/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py index 3fd2e8808e7..8c1e9e48661 100644 --- a/tests/unit/router_strategy/test_llm_v2.py +++ b/tests/unit/router_strategy/test_llm_v2.py @@ -4,11 +4,12 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest -import litellm from pydantic import ValidationError +import litellm from litellm import ModelResponse, Router from litellm.caching.dual_cache import DualCache +from litellm.router_strategy.auto_router.auto_router_model_naming import strategy_router_dependencies from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier from litellm.router_strategy.complexity_router.fuse_presets import get_fuse_presets @@ -21,7 +22,6 @@ from litellm.router_strategy.complexity_router.llm_v2 import ( LLMV2Verdict, llm_v2_response_format, ) -from litellm.router_utils.auto_router_model_naming import strategy_router_dependencies from litellm.types.llms.openai import ResponsesAPIResponse diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 0d39016ea79..2e1e3acaabb 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -25,6 +25,7 @@ from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callba from litellm.router_utils.cooldown_handlers import ( _get_deployment_cooldown_policy, _has_explicit_allowed_fails_policy_for_exception, + _has_fallback_group, _increment_allowed_fails, _is_cooldown_required, _resolve_allowed_fails_from_policy, @@ -3180,6 +3181,34 @@ async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_poli assert len(route.calls) == outcomes.count(litellm.Timeout) +@pytest.mark.parametrize( + ("fallbacks", "expected"), + [ + ([{"test-model": "backup-model"}], True), + ([{"*": "backup-model"}], True), + ([{"test-model": ["backup-model"]}], True), + ([{"test-model": 5}, {"*": "backup-model"}], True), + ([{"test-model": 5}], False), + ([{"other-model": ["backup-model"]}], False), + ], + ids=[ + "string_target", + "string_catch_all_from_default_fallbacks", + "list_target", + "unparseable_entry_does_not_hide_the_catch_all", + "unparseable_entry_alone", + "another_groups_chain", + ], +) +def test_has_fallback_group_reads_string_targets_and_skips_unparseable_entries( + fallbacks: Sequence[Mapping[str, object]], expected: bool +) -> None: + router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT]) + router.fallbacks = list(fallbacks) + + assert _has_fallback_group(router, "test-model") is expected + + async def _group_call_outcome(router: Router) -> str: try: response: Final = await router.acompletion( diff --git a/tests/unit/rust_bridge/host/test_secret_manager.py b/tests/unit/rust_bridge/host/test_secret_manager.py index 9d876124ffd..1f8acc5792e 100644 --- a/tests/unit/rust_bridge/host/test_secret_manager.py +++ b/tests/unit/rust_bridge/host/test_secret_manager.py @@ -17,6 +17,7 @@ from botocore.credentials import Credentials from pydantic import JsonValue import litellm +from litellm.rust_bridge import _native as native from litellm.rust_bridge.host.secret_manager import capture_secret_manager, native_secret_manager_config from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager @@ -67,7 +68,6 @@ def _vault_body(value: str) -> dict[str, object]: def _native_handle(client: object) -> "_SecretManagerRuntime": - native: Final = pytest.importorskip("litellm.rust_bridge._native") handle: Final = native._SecretManagerRuntime.from_client(client) assert handle is not None return handle @@ -75,7 +75,6 @@ def _native_handle(client: object) -> "_SecretManagerRuntime": @pytest.mark.parametrize("system", ("aws_secret_manager", "hashicorp_vault", "cyberark")) async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(system: str) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: environment: Final = { "AWS_REGION_NAME": "us-east-1", @@ -109,7 +108,6 @@ async def test_python_native_handle_reuses_backend_across_sync_and_async_reads(s def test_shared_initializer_captures_credentials_and_tracks_instance_settings(monkeypatch: pytest.MonkeyPatch) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": "native-value"}) server.expected_requests = 2 @@ -141,7 +139,6 @@ def test_shared_initializer_captures_credentials_and_tracks_instance_settings(mo def test_configuration_replacement_rebuilds_without_invalidating_existing_handles( monkeypatch: pytest.MonkeyPatch, ) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as first_server, recording_service() as second_server: first_server.default_response = ResponseSpec(body=_vault_body("first")) second_server.default_response = ResponseSpec(body=_vault_body("second")) @@ -158,7 +155,6 @@ def test_configuration_replacement_rebuilds_without_invalidating_existing_handle def test_custom_subclass_is_not_captured_for_a_native_handle() -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") class CustomManager(AWSSecretsManagerV2): def sync_read_secret(self, secret_name: str, primary_secret_name: str | None = None) -> str: @@ -171,19 +167,16 @@ def test_custom_subclass_is_not_captured_for_a_native_handle() -> None: @pytest.mark.parametrize("system", ("google_secret_manager", "hashicorp_vault", "cyberark")) def test_enterprise_backends_cannot_initialize_without_entitlement(system: str) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with pytest.raises(ValueError, match=r"[Ee]nterprise|[Pp]remium"): native._SecretManagerRuntime.from_config(system, {"CYBERARK_API_KEY": "key"}) def test_azure_factory_rejects_unencrypted_vault_endpoints() -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with pytest.raises(ValueError, match="https"): native._SecretManagerRuntime.from_config("azure_key_vault", {"AZURE_KEY_VAULT_URI": "http://127.0.0.1:1"}) def test_direct_builtin_constructor_can_use_native_without_registration(monkeypatch: pytest.MonkeyPatch) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body=_vault_body("native-value")) manager: Final = _vault(monkeypatch, server.base_url) @@ -193,12 +186,10 @@ def test_direct_builtin_constructor_can_use_native_without_registration(monkeypa def test_uncaptured_clients_have_no_native_handle() -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") assert native._SecretManagerRuntime.from_client(SimpleNamespace()) is None def test_native_failure_is_not_replayed_in_python(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(status=403, body={"errors": ["denied"]}) server.expected_requests = 1 @@ -267,7 +258,6 @@ def test_same_named_custom_client_is_not_captured() -> None: @pytest.mark.parametrize("value", ("text", "", True, False, 42, 2**100, [1, "two"], {"nested": True}, None)) async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.MonkeyPatch, value: JsonValue) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": json.dumps({"KEY": value})}) server.expected_requests = 3 @@ -290,7 +280,6 @@ async def test_aws_primary_values_match_python_handler(monkeypatch: pytest.Monke @pytest.mark.parametrize("primary", (None, "primary")) async def test_aws_primary_values_match_python_reads(monkeypatch: pytest.MonkeyPatch, primary: str | None) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": '{"KEY":"value"}'}) server.expected_requests = 3 @@ -343,7 +332,6 @@ def test_aws_absence_and_failed_reads_match_python_without_environment_fallback( async def test_aws_primary_json_errors_preserve_python_exception_details( monkeypatch: pytest.MonkeyPatch, document: str ) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": document}) server.expected_requests = 3 @@ -370,7 +358,6 @@ async def test_aws_primary_json_errors_preserve_python_exception_details( async def test_public_aws_reads_preserve_coroutines_and_per_call_credentials( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as initial, recording_service() as selected: initial.expected_requests = 0 selected.expected_requests = 2 @@ -421,7 +408,6 @@ async def test_public_aws_reads_preserve_coroutines_and_per_call_credentials( async def test_public_aws_primary_reads_ignore_operation_overrides(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server, recording_service() as unused: server.default_response = ResponseSpec(body={"SecretString": '{"KEY":true}'}) server.expected_requests = 2 @@ -440,7 +426,6 @@ async def test_public_aws_primary_reads_ignore_operation_overrides(monkeypatch: async def test_public_aws_bootstrap_names_only_bypass_sync_reads(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": "remote-access"}) server.expected_requests = 1 @@ -457,7 +442,6 @@ async def test_public_aws_bootstrap_names_only_bypass_sync_reads(monkeypatch: py async def test_public_aws_read_timeouts_follow_the_python_http_handler( monkeypatch: pytest.MonkeyPatch, timeout: float | httpx.Timeout ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body={"SecretString": "too-late"}, delay=0.25) server.expected_requests = 2 @@ -470,7 +454,6 @@ async def test_public_aws_read_timeouts_follow_the_python_http_handler( async def test_public_vault_reads_keep_overrides_cache_and_coroutines(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body=_vault_body("value")) server.expected_requests = 2 @@ -490,7 +473,6 @@ async def test_public_vault_reads_keep_overrides_cache_and_coroutines(monkeypatc async def test_public_vault_failed_reads_return_none_without_replay( monkeypatch: pytest.MonkeyPatch, status: int ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(status=status, body={"errors": ["unavailable"]}) server.expected_requests = 2 @@ -502,7 +484,6 @@ async def test_public_vault_failed_reads_return_none_without_replay( async def test_public_cyberark_reads_reuse_authentication_and_cached_values( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") from litellm.proxy import proxy_server from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager @@ -536,7 +517,6 @@ def test_public_aws_bootstrap_read_does_not_hit_the_provider(monkeypatch: pytest def test_public_vault_prefix_overrides_match_python_string_conversion( monkeypatch: pytest.MonkeyPatch, override: str | int ) -> None: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("HCP_VAULT_PATH_PREFIX", "default-prefix") with recording_service() as server: server.default_response = ResponseSpec(body=_vault_body("value")) @@ -562,7 +542,6 @@ def _cyberark(monkeypatch: pytest.MonkeyPatch, address: str) -> CyberArkSecretMa async def test_public_cyberark_writes_and_deletes_share_the_read_cache( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for body in (b"token", b"old", {}, {}, b"provider-after-delete"): server.enqueue(ResponseSpec(body=body)) @@ -609,7 +588,6 @@ async def test_public_cyberark_write_errors_match_python_without_http_retries( status: int, authentication: bool, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: responses: Final = ( (ResponseSpec(status=status, body={}),) * 2 @@ -635,7 +613,6 @@ async def test_public_cyberark_write_errors_match_python_without_http_retries( async def test_public_cyberark_write_recovers_from_initial_policy_authentication_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.enqueue(ResponseSpec(status=401, body={})) server.enqueue(ResponseSpec(body=b"token")) @@ -658,7 +635,6 @@ async def test_public_cyberark_write_rejects_unsafe_names_before_authentication( monkeypatch: pytest.MonkeyPatch, name: str, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.expected_requests = 0 manager: Final = _cyberark(monkeypatch, server.base_url) @@ -673,7 +649,6 @@ async def test_public_cyberark_rotation_returns_the_write_response_and_retains_o monkeypatch: pytest.MonkeyPatch, same_name: bool, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for body in (b"token", b"old-value", {}, {}): server.enqueue(ResponseSpec(body=body)) @@ -696,7 +671,6 @@ async def test_public_native_cyberark_rotation_requires_a_fresh_matching_replace monkeypatch: pytest.MonkeyPatch, replacement: bytes | None, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for body in (b"token", b"old-value", {}, {}): server.enqueue(ResponseSpec(body=body)) @@ -720,7 +694,6 @@ async def test_public_native_cyberark_rotation_requires_a_fresh_matching_replace async def test_public_cyberark_cached_authentication_does_not_retry_denied_reads( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.enqueue(ResponseSpec(body=b"token")) server.enqueue(ResponseSpec(body=b"value")) @@ -734,7 +707,6 @@ async def test_public_cyberark_cached_authentication_does_not_retry_denied_reads async def test_cyberark_handler_errors_match_python_after_cached_authentication_is_denied( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for response in ( ResponseSpec(body=b"token"), @@ -757,7 +729,6 @@ async def test_cyberark_handler_errors_match_python_after_cached_authentication_ async def test_public_cyberark_connection_errors_match_python( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.expected_requests = 0 address: Final = server.base_url @@ -769,7 +740,6 @@ async def test_public_cyberark_connection_errors_match_python( async def test_public_native_cyberark_rotation_stops_after_a_failed_write(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: for body in (b"token", b"old-value", {}): server.enqueue(ResponseSpec(body=body)) @@ -788,7 +758,6 @@ async def test_public_vault_writes_preserve_complete_responses_and_request_field monkeypatch: pytest.MonkeyPatch, description: str | None, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: response: Final = { "request_id": "test-request", @@ -829,7 +798,6 @@ async def test_public_vault_mutation_http_errors_match_python_without_retry( operation: str, status: int, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(status=status, body={"errors": ["denied"]}) server.expected_requests = 2 @@ -856,7 +824,6 @@ async def test_public_vault_write_response_conversion_matches_python( monkeypatch: pytest.MonkeyPatch, body: bytes, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body=body) server.expected_requests = 2 @@ -870,7 +837,6 @@ async def test_public_vault_write_response_conversion_matches_python( async def test_public_vault_deletion_invalidates_cached_fields( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.enqueue(ResponseSpec(body=_vault_body("old"))) server.enqueue(ResponseSpec(status=204, body=b"")) @@ -893,7 +859,6 @@ async def test_public_vault_rotation_preserves_response_and_best_effort_deletion same_name: bool, delete_status: int, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: response: Final = {"request_id": "write-id", "data": {"version": 3}, "extra": [1, 2]} server.enqueue(ResponseSpec(body=b"current-existence-is-status-only")) @@ -924,7 +889,6 @@ async def test_public_vault_rotation_failure_messages_and_request_counts_match_p stage: str, status: int, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: before: Final = ( () @@ -960,7 +924,6 @@ async def test_public_vault_rotation_mismatches_do_not_delete_the_old_alias( monkeypatch: pytest.MonkeyPatch, value: JsonValue, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: responses: Final = ( ResponseSpec(body=_vault_body("old")), @@ -983,7 +946,6 @@ async def test_public_vault_mutation_timeouts_match_python( monkeypatch: pytest.MonkeyPatch, operation: str, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.default_response = ResponseSpec(body=_vault_body("value"), delay=0.25) server.expected_requests = 2 @@ -1021,7 +983,6 @@ async def test_public_vault_unsafe_names_fail_before_authentication( monkeypatch: pytest.MonkeyPatch, operation: str, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.expected_requests = 0 manager: Final = _vault(monkeypatch, server.base_url) @@ -1038,7 +999,6 @@ async def test_public_vault_authentication_errors_match_python( monkeypatch: pytest.MonkeyPatch, operation: str, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("HCP_VAULT_APPROLE_ROLE_ID", "role") monkeypatch.setenv("HCP_VAULT_APPROLE_SECRET_ID", "secret-id") with recording_service() as server: @@ -1064,7 +1024,6 @@ async def test_public_vault_authentication_errors_match_python( async def test_public_vault_rotation_stops_on_a_success_response_containing_an_error( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: result: Final = {"status": "error", "message": "write rejected", "extra": 2**100} for response in (ResponseSpec(body=_vault_body("old")), ResponseSpec(body=result)) * 2: @@ -1078,7 +1037,6 @@ async def test_public_vault_rotation_stops_on_a_success_response_containing_an_e async def test_public_native_vault_write_invalidates_stale_cached_values(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.enqueue(ResponseSpec(body=_vault_body("old"))) server.enqueue(ResponseSpec(body={"data": {"version": 2}})) @@ -1094,7 +1052,6 @@ async def test_public_native_vault_write_invalidates_stale_cached_values(monkeyp async def test_public_native_vault_write_rejects_description_overwriting_the_secret( monkeypatch: pytest.MonkeyPatch, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: server.expected_requests = 0 handle: Final = _native_handle(_vault(monkeypatch, server.base_url)) @@ -1111,7 +1068,6 @@ async def test_public_vault_rotation_preserves_malformed_verification_errors( monkeypatch: pytest.MonkeyPatch, body: JsonValue, ) -> None: - pytest.importorskip("litellm.rust_bridge._native") with recording_service() as server: responses: Final = ( ResponseSpec(body=_vault_body("old")), diff --git a/tests/unit/rust_bridge/messages/test_secrets.py b/tests/unit/rust_bridge/messages/test_secrets.py index f6439753809..64559d63193 100644 --- a/tests/unit/rust_bridge/messages/test_secrets.py +++ b/tests/unit/rust_bridge/messages/test_secrets.py @@ -17,8 +17,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE from tests.test_litellm_rust.support.recording_server import ResponseSpec, recording_service -pytest.importorskip("litellm.rust_bridge._native") - pytestmark = pytest.mark.usefixtures("local_model_cost_map") diff --git a/tests/unit/rust_bridge/ocr/test_secrets.py b/tests/unit/rust_bridge/ocr/test_secrets.py index 33ded916888..4ea6e09f885 100644 --- a/tests/unit/rust_bridge/ocr/test_secrets.py +++ b/tests/unit/rust_bridge/ocr/test_secrets.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge import _native as native from litellm.rust_bridge.host import settings from litellm.rust_bridge.ocr.entrypoints import NATIVE_AOCR, NATIVE_OCR from litellm.rust_bridge.public_call import NativeCall @@ -19,8 +20,6 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem from tests.test_litellm_rust.ocr.data import OCR_DOCUMENT, OCR_MODEL, OCR_RESPONSE from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec, recording_service -native: Final = pytest.importorskip("litellm.rust_bridge._native") - AccessMode: TypeAlias = Literal["read_only", "write_only", "read_and_write"] diff --git a/tests/unit/rust_bridge/test_bindings.py b/tests/unit/rust_bridge/test_bindings.py index 39a157220dc..086be65ee26 100644 --- a/tests/unit/rust_bridge/test_bindings.py +++ b/tests/unit/rust_bridge/test_bindings.py @@ -6,7 +6,7 @@ from typing import Final, cast # noqa: TID251 # narrows the async binding's un import pytest -from litellm.rust_bridge import bindings +from litellm.rust_bridge import _native, bindings from litellm.rust_bridge.chat_completions import entrypoints as chat_completions from litellm.rust_bridge.embeddings import entrypoints as embeddings from litellm.rust_bridge.execution import Complete, Step @@ -17,12 +17,15 @@ from litellm.rust_bridge.responses import entrypoints as responses from litellm.rust_bridge.transcription import native as transcription -def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None: - native = SimpleNamespace(route=lambda: "native") - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) - binding: bindings.NativeBinding[object] = bindings.NativeBinding("route", validate=lambda value: value) +def test_binding_distinguishes_disable_from_reset() -> None: + def route() -> str: + return "native" - assert binding.load() is native.route + binding: bindings.NativeBinding[object] = bindings.NativeBinding( + "route", validate=lambda value: value, source=SimpleNamespace(route=route) + ) + + assert binding.load() is route binding.override(None) assert binding.load() is None @@ -32,16 +35,15 @@ def test_binding_distinguishes_disable_from_reset(monkeypatch) -> None: assert binding.load() is replacement binding.reset() - assert binding.load() is native.route + assert binding.load() is route @pytest.mark.parametrize(("value", "expected"), ((3, 3), ("invalid", None), (None, None))) -def test_binding_validates_native_attribute( - monkeypatch: pytest.MonkeyPatch, value: object, expected: int | None -) -> None: +def test_binding_validates_native_attribute(value: object, expected: int | None) -> None: native: Final = SimpleNamespace(route=value) - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) - binding: Final = bindings.NativeBinding("route", validate=lambda item: item if isinstance(item, int) else None) + binding: Final = bindings.NativeBinding( + "route", validate=lambda item: item if isinstance(item, int) else None, source=native + ) assert binding.load() == expected @@ -72,28 +74,33 @@ ROUTE_BINDINGS: Final = tuple((attribute, binding) for attribute, binding, _ in @pytest.mark.parametrize( ("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS] ) -def test_route_bindings_reject_non_callable_native_attributes( - monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object] +def test_route_bindings_resolve_the_packaged_native_route( + attribute: str, route_binding: bindings.NativeBinding[object] ) -> None: - monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: "not callable"})) - route_binding.reset() - assert route_binding.load() is None route_binding.reset() + assert attribute in dir(_native) + assert callable(route_binding.load()) + + +@pytest.mark.parametrize( + ("attribute", "route_binding"), ROUTE_BINDINGS, ids=[attribute for attribute, _ in ROUTE_BINDINGS] +) +def test_route_bindings_reject_non_callable_native_attributes( + attribute: str, route_binding: bindings.NativeBinding[object] +) -> None: + assert route_binding.with_source(SimpleNamespace(**{attribute: "not callable"})).load() is None @pytest.mark.parametrize( ("attribute", "route_binding"), TRANSCRIPTION_BINDINGS, ids=[attribute for attribute, _ in TRANSCRIPTION_BINDINGS] ) def test_transcription_bindings_expose_the_native_callable_itself( - monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object] + attribute: str, route_binding: bindings.NativeBinding[object] ) -> None: def native_route() -> None: pass - monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route})) - route_binding.reset() - assert route_binding.load() is native_route - route_binding.reset() + assert route_binding.with_source(SimpleNamespace(**{attribute: native_route})).load() is native_route @dataclass(frozen=True, slots=True) @@ -116,7 +123,6 @@ class _EchoingExecution: ids=[attribute for attribute, _, _ in EXECUTION_BINDINGS], ) def test_execution_bindings_start_the_native_execution_and_return_its_result( - monkeypatch: pytest.MonkeyPatch, attribute: str, route_binding: bindings.NativeBinding[object], asynchronous: bool, @@ -128,9 +134,7 @@ def test_execution_bindings_start_the_native_execution_and_return_its_result( received.append(incoming) return _EchoingExecution() - monkeypatch.setattr(bindings, "get_native_bridge", lambda: SimpleNamespace(**{attribute: native_route})) - route_binding.reset() - bound: Final = route_binding.load() + bound: Final = route_binding.with_source(SimpleNamespace(**{attribute: native_route})).load() assert callable(bound) outcome: Final = bound(call) @@ -139,7 +143,6 @@ def test_execution_bindings_start_the_native_execution_and_return_its_result( if asynchronous else outcome ) - route_binding.reset() assert received == [call] assert result == ("completed with", None) diff --git a/tests/unit/rust_bridge/test_fork_guard.py b/tests/unit/rust_bridge/test_fork_guard.py index 88ae017ec39..aad50437f03 100644 --- a/tests/unit/rust_bridge/test_fork_guard.py +++ b/tests/unit/rust_bridge/test_fork_guard.py @@ -1,36 +1,21 @@ -from types import SimpleNamespace - import pytest from litellm.rust_bridge import fork_guard -def _reserve_with(monkeypatch: pytest.MonkeyPatch, native: object) -> None: - monkeypatch.setattr(fork_guard, "get_native_bridge", lambda: native) - fork_guard.reserve_process_for_forking("the gunicorn master") - - -def test_missing_extension_has_nothing_to_reserve(monkeypatch: pytest.MonkeyPatch) -> None: - assert _reserve_with(monkeypatch, None) is None - - -def test_extension_built_before_reservation_existed_passes(monkeypatch: pytest.MonkeyPatch) -> None: - assert _reserve_with(monkeypatch, SimpleNamespace()) is None - - -def test_unused_extension_is_reserved(monkeypatch: pytest.MonkeyPatch) -> None: +def test_unused_extension_is_reserved() -> None: calls: list[None] = [] - _reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=lambda: calls.append(None))) + fork_guard.reserve_process_for_forking("the gunicorn master", reserve=lambda: calls.append(None)) assert calls == [None] -def test_used_extension_refuses_and_names_the_place(monkeypatch: pytest.MonkeyPatch) -> None: +def test_used_extension_refuses_and_names_the_place() -> None: def reserve() -> None: raise RuntimeError("the native runtime already started in this process") with pytest.raises(fork_guard.NativeStateStartedBeforeFork, match="the gunicorn master") as raised: - _reserve_with(monkeypatch, SimpleNamespace(reserve_process_for_forking=reserve)) + fork_guard.reserve_process_for_forking("the gunicorn master", reserve=reserve) assert isinstance(raised.value.__cause__, RuntimeError) diff --git a/tests/unit/rust_bridge/test_logger.py b/tests/unit/rust_bridge/test_logger.py index 5a7b6c270a1..4745aca9d50 100644 --- a/tests/unit/rust_bridge/test_logger.py +++ b/tests/unit/rust_bridge/test_logger.py @@ -80,7 +80,6 @@ def test_native_logging_observes_level_changes(caplog: pytest.LogCaptureFixture) ), ) def test_native_credential_patterns_match_python(text: str) -> None: - pytest.importorskip("litellm.rust_bridge._native") from litellm.rust_bridge._native import NativeDiagnosticProcessor processor: Final = NativeDiagnosticProcessor(MINIMUM_CUSTOM_KEY_LENGTH) @@ -91,7 +90,6 @@ def test_native_credential_patterns_match_python(text: str) -> None: def test_native_client_redaction_matches_python() -> None: - pytest.importorskip("litellm.rust_bridge._native") from litellm.rust_bridge._native import NativeDiagnosticProcessor text: Final = "error at /etc/secrets/config on db.internal\nTraceback (most recent call last):\nsecret" @@ -100,7 +98,6 @@ def test_native_client_redaction_matches_python() -> None: def test_native_diagnostic_batch_matches_python() -> None: - pytest.importorskip("litellm.rust_bridge._native") from litellm.rust_bridge._native import NativeDiagnosticProcessor message: Final = "é" * 110 + "sk-" + "q" * 48 + "界" * 1000 diff --git a/tests/unit/rust_bridge/test_runtime.py b/tests/unit/rust_bridge/test_runtime.py index 95fa3e2cc7b..c0c416a62c9 100644 --- a/tests/unit/rust_bridge/test_runtime.py +++ b/tests/unit/rust_bridge/test_runtime.py @@ -1,7 +1,6 @@ from __future__ import annotations from collections.abc import Callable, Generator -from types import SimpleNamespace from typing import Final, Protocol import pytest @@ -10,26 +9,14 @@ from litellm.exceptions import APIError from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.rust_bridge import bindings, configuration, runtime +from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError from litellm.rust_bridge.catalog import Route, RouteContext, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.rust_bridge.execution import AsyncStream, Complete, StreamChunk, SyncStream -class RustBridgeDeclined(Exception): - pass - - -class RustUpstreamError(Exception): - pass - - @pytest.fixture(autouse=True) -def native_exceptions(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: - native: Final = SimpleNamespace( - RustBridgeDeclined=RustBridgeDeclined, - RustUpstreamError=RustUpstreamError, - ) - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) +def rust_configuration(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: monkeypatch.delenv("LITELLM_RUST", raising=False) configuration.reset_rust_configuration() yield diff --git a/tests/unit/rust_bridge/test_token_counter.py b/tests/unit/rust_bridge/test_token_counter.py index 3af7127a9b3..3513a0ef44c 100644 --- a/tests/unit/rust_bridge/test_token_counter.py +++ b/tests/unit/rust_bridge/test_token_counter.py @@ -16,6 +16,7 @@ import litellm from litellm.constants import TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding from litellm.proxy.spend_tracking.input_tokens import count_input_tokens_for_model +from litellm.rust_bridge import _native as native from litellm.rust_bridge import token_counter as bridge from litellm.rust_bridge import tokenizer as tokenizer_dispatch from litellm.rust_bridge._native import Tokenizer @@ -313,7 +314,6 @@ PARITY_MODELS: Final[tuple[tuple[str, bridge.RustTokenizer], ...]] = ( async def test_native_count_matches_python_budget_counter( request_body: dict[str, object], model: str, tokenizer: bridge.RustTokenizer ) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") body: Final = json.dumps(request_body).replace(MODEL, model) parsed: Final = json.loads(body) @@ -328,7 +328,6 @@ async def test_tiktoken_counts_long_text_exactly_where_python_chunks( model: str, tokenizer: bridge.RustTokenizer ) -> None: """Python encodes tiktoken text in fixed-size chunks (drift of up to one token per chunk boundary); Rust does not.""" - native: Final = pytest.importorskip("litellm.rust_bridge._native") text: Final = "x " * 20_000 body: Final = {"model": model, "messages": [{"role": "user", "content": text}]} encoding: Final = Tokenizer.from_tiktoken(tokenizer) @@ -362,7 +361,6 @@ DECLINED_REQUESTS: Final[tuple[dict[str, object], ...]] = ( async def test_native_declines_shapes_python_prices_differently( request_body: dict[str, object], tokenizer: bridge.RustTokenizer ) -> None: - native: Final = pytest.importorskip("litellm.rust_bridge._native") raw, _ = _counted(request_body, MODEL_BY_TOKENIZER[tokenizer]) with pytest.raises(native.RustBridgeDeclined): diff --git a/tests/unit/sandbox/test_e2b_sandbox.py b/tests/unit/sandbox/test_e2b_sandbox.py index e01b9120416..e58392ca266 100644 --- a/tests/unit/sandbox/test_e2b_sandbox.py +++ b/tests/unit/sandbox/test_e2b_sandbox.py @@ -7,6 +7,7 @@ integration tests live in tests/integration/sandbox/test_e2b_sandbox.py. """ import json +from typing import Final import httpx import pytest @@ -19,19 +20,18 @@ from litellm.llms.e2b.sandbox.transformation import ( ) +def _line_stream(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) + + class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): + def __init__(self, *, json_data: object = None, status_code: int = 200) -> None: self._json = json_data - self._lines = lines or [] self.status_code = status_code def json(self): return self._json - async def aiter_lines(self): - for line in self._lines: - yield line - class FakeHTTPClient: """Records outbound requests and returns canned responses keyed by URL.""" @@ -61,7 +61,7 @@ class FakeHTTPClient: if url.endswith("/execute"): if self.execute_raises is not None: raise self.execute_raises - return FakeResponse(lines=self.execute_lines) + return _line_stream(self.execute_lines) raise AssertionError(f"unexpected POST {url}") async def delete(self, url, headers=None, **kwargs): @@ -249,6 +249,20 @@ def test_parse_lines_skips_non_json_lines(): assert result.error is None +@pytest.mark.asyncio +async def test_run_code_keeps_unicode_line_separators_inside_stdout() -> None: + separated_text: Final = "before\u2028between\u2029after\n" + client: Final = FakeHTTPClient( + execute_lines=[json.dumps({"type": "stdout", "text": separated_text, "timestamp": 1}, ensure_ascii=False)] + ) + handle: Final = ContainerHandle(id="sbx_sep", provider="e2b", domain="e2b.app") + handle._hidden_params = {"envd_access_token": "tok"} + + result: Final = await E2BSandboxConfig().arun_code(container=handle, code="print(text)", client=client) + + assert result.stdout == separated_text + + @pytest.mark.asyncio async def test_run_code_aborts_on_output_over_cap(): big_line = "x" * (MAX_OUTPUT_BYTES + 1) diff --git a/tests/unit/sandbox/test_opensandbox_sandbox.py b/tests/unit/sandbox/test_opensandbox_sandbox.py index 2928dea100e..36ae848786c 100644 --- a/tests/unit/sandbox/test_opensandbox_sandbox.py +++ b/tests/unit/sandbox/test_opensandbox_sandbox.py @@ -27,10 +27,13 @@ def sse(data): return f"data: {json.dumps(data)}" +def _line_stream(lines: list[str]) -> httpx.Response: + return httpx.Response(200, stream=httpx.ByteStream("".join(f"{line}\n" for line in lines).encode())) + + class FakeResponse: - def __init__(self, *, json_data=None, lines=None, status_code=200): + def __init__(self, *, json_data: object = None, status_code: int = 200) -> None: self._json = json_data - self._lines = lines or [] self.status_code = status_code def json(self): @@ -40,10 +43,6 @@ class FakeResponse: if self.status_code >= 400: raise http_status_error(self.status_code) - async def aiter_lines(self): - for line in self._lines: - yield line - class FakeHTTPClient: def __init__( @@ -93,7 +92,7 @@ class FakeHTTPClient: if url.endswith("/code"): if self.execute_raises is not None: raise self.execute_raises - return FakeResponse(lines=self.execute_lines) + return _line_stream(self.execute_lines) raise AssertionError(f"unexpected POST {url}") async def get(self, url, headers=None, params=None, **kwargs): diff --git a/tests/unit/skills/test_skills_data/slack-gif-creator/SKILL.md b/tests/unit/skills/test_skills_data/slack-gif-creator/SKILL.md index 3cae971b731..d2cdfd17bde 100644 --- a/tests/unit/skills/test_skills_data/slack-gif-creator/SKILL.md +++ b/tests/unit/skills/test_skills_data/slack-gif-creator/SKILL.md @@ -1,6 +1,8 @@ --- name: slack-gif-creator description: Knowledge and utilities for creating animated GIFs optimized for Slack. Provides constraints, validation tools, and animation concepts. Use when users request animated GIFs for Slack like "make me a GIF of X doing Y for Slack." +metadata: + internal: true --- # Slack GIF Creator diff --git a/tests/unit/skills/test_skills_data/test-delete-skill/SKILL.md b/tests/unit/skills/test_skills_data/test-delete-skill/SKILL.md index ba3fda9ab7b..ef1fd1a7b24 100644 --- a/tests/unit/skills/test_skills_data/test-delete-skill/SKILL.md +++ b/tests/unit/skills/test_skills_data/test-delete-skill/SKILL.md @@ -1,6 +1,8 @@ --- name: test-delete-skill description: A test skill created specifically for deletion testing +metadata: + internal: true --- # Test Delete Skill diff --git a/tests/unit/skills/test_skills_data/test-skill-litellm/SKILL.md b/tests/unit/skills/test_skills_data/test-skill-litellm/SKILL.md index ec83b580f0e..b3d54ccdf3c 100644 --- a/tests/unit/skills/test_skills_data/test-skill-litellm/SKILL.md +++ b/tests/unit/skills/test_skills_data/test-skill-litellm/SKILL.md @@ -1,6 +1,8 @@ --- name: test-skill-litellm description: A test skill created by LiteLLM automated tests +metadata: + internal: true --- # Test Skill diff --git a/tests/unit/test_audio_transcription_rust_bridge.py b/tests/unit/test_audio_transcription_rust_bridge.py index e3a6781dee7..715423c59d9 100644 --- a/tests/unit/test_audio_transcription_rust_bridge.py +++ b/tests/unit/test_audio_transcription_rust_bridge.py @@ -1,14 +1,14 @@ from __future__ import annotations from collections.abc import Generator -from types import SimpleNamespace from typing import Final import pytest import litellm from litellm.llms.bedrock.audio_transcription import BedrockAudioTranscriptionRustDispatch -from litellm.rust_bridge import bindings, configuration +from litellm.rust_bridge import configuration +from litellm.rust_bridge._native import RustBridgeDeclined, RustUpstreamError from litellm.rust_bridge.public_call import NativeCall from litellm.rust_bridge.transcription.native import NATIVE_ATRANSCRIPTION, NATIVE_TRANSCRIPTION @@ -28,18 +28,8 @@ TRANSCRIPTION_FIELDS: Final = frozenset( ) -class RustBridgeDeclined(Exception): - pass - - -class RustUpstreamError(Exception): - pass - - @pytest.fixture(autouse=True) def isolated_bridge(monkeypatch: pytest.MonkeyPatch) -> Generator[None]: - native: Final = SimpleNamespace(RustBridgeDeclined=RustBridgeDeclined, RustUpstreamError=RustUpstreamError) - monkeypatch.setattr(bindings, "get_native_bridge", lambda: native) monkeypatch.delenv("LITELLM_RUST", raising=False) configuration.reset_rust_configuration() yield diff --git a/tests/unit/test_groq_streaming_encoding.py b/tests/unit/test_groq_streaming_encoding.py index ffe87fb9890..37f05272f67 100644 --- a/tests/unit/test_groq_streaming_encoding.py +++ b/tests/unit/test_groq_streaming_encoding.py @@ -6,148 +6,58 @@ UTF-8 encoded content in streaming responses, specifically fixing the ASCII encoding error described in issue #12660. """ -import asyncio -from unittest.mock import AsyncMock, Mock +import json +from typing import Final +from unittest.mock import Mock +import httpx import pytest +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.openai_like.chat.handler import make_call, make_sync_call - -class MockResponse: - """Mock httpx response for testing UTF-8 handling.""" - - def __init__(self, test_content: str): - self.test_content = test_content - self.status_code = 200 - - def iter_text(self, encoding="utf-8"): - """Mock iter_text that yields content with the specified encoding.""" - yield self.test_content - - async def aiter_text(self, encoding="utf-8"): - """Mock aiter_text that yields content with the specified encoding.""" - yield self.test_content - - def iter_lines(self): - """Mock iter_lines method for synchronous streaming.""" - yield self.test_content - - async def aiter_lines(self): - """Mock aiter_lines method for asynchronous streaming.""" - yield self.test_content - - def json(self): - return {"choices": [{"delta": {"content": "test"}}]} +_UNICODE_TEXTS: Final = ( + "The symbol µ represents micro", + "Testing © character", + "Testing ™ character", + "Testing € character", + "Testing 北京 character", + "Testing 🚀 character", + "Testing Ñoño character", +) -class MockSyncClient: - """Mock synchronous HTTP client for testing.""" - - def __init__(self, response_content: str): - self.response_content = response_content - - def post(self, *args, **kwargs): - return MockResponse(self.response_content) +def _utf8_stream_transport(text: str) -> httpx.MockTransport: + chunk: Final = {"choices": [{"index": 0, "delta": {"content": text}}]} + body: Final = f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode() + return httpx.MockTransport(lambda request: httpx.Response(200, stream=httpx.ByteStream(body))) -class MockAsyncClient: - """Mock asynchronous HTTP client for testing.""" - - def __init__(self, response_content: str): - self.response_content = response_content - - async def post(self, *args, **kwargs): - return MockResponse(self.response_content) +def _call_args() -> dict[str, object]: + return { + "api_base": "https://test.com/v1/chat/completions", + "headers": {"Authorization": "Bearer test"}, + "data": '{"model": "test", "messages": []}', + "model": "test-model", + "messages": [], + "logging_obj": Mock(), + } -def test_utf8_streaming_sync(): - """Test that synchronous streaming handles UTF-8 characters correctly.""" - # Content with the µ character that was causing issues - test_content = ( - 'data: {"choices":[{"delta":{"content":"The symbol µ represents micro"}}]}\n\n' +@pytest.mark.parametrize("text", _UNICODE_TEXTS) +def test_utf8_streaming_sync(text: str) -> None: + completion_stream: Final = make_sync_call( + client=HTTPHandler(client=httpx.Client(transport=_utf8_stream_transport(text))), **_call_args() ) - mock_client = MockSyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error - completion_stream = make_sync_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - # Verify we can iterate through the stream without encoding errors - assert completion_stream is not None + assert "".join(chunk["text"] for chunk in completion_stream) == text @pytest.mark.asyncio -async def test_utf8_streaming_async(): - """Test that asynchronous streaming handles UTF-8 characters correctly.""" - # Content with the µ character that was causing issues - test_content = ( - 'data: {"choices":[{"delta":{"content":"The symbol µ represents micro"}}]}\n\n' +@pytest.mark.parametrize("text", _UNICODE_TEXTS) +async def test_utf8_streaming_async(text: str) -> None: + completion_stream: Final = await make_call( + client=AsyncHTTPHandler(transport=_utf8_stream_transport(text)), **_call_args() ) - mock_client = MockAsyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error - completion_stream = await make_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - # Verify we can iterate through the stream without encoding errors - assert completion_stream is not None - - -def test_various_unicode_characters(): - """Test streaming with various Unicode characters that could cause issues.""" - unicode_test_cases = [ - "µ", # Micro symbol (the original issue) - "©", # Copyright symbol - "™", # Trademark symbol - "€", # Euro symbol - "北京", # Chinese characters - "🚀", # Emoji - "Ñoño", # Spanish characters with tildes - ] - - for unicode_char in unicode_test_cases: - test_content = f'data: {{"choices":[{{"delta":{{"content":"Testing {unicode_char} character"}}}}]}}\n\n' - - mock_client = MockSyncClient(test_content) - mock_logging = Mock() - - # This should not raise an ASCII encoding error for any Unicode character - completion_stream = make_sync_call( - client=mock_client, - api_base="https://test.com/v1/chat/completions", - headers={"Authorization": "Bearer test"}, - data='{"model": "test", "messages": []}', - model="test-model", - messages=[], - logging_obj=mock_logging, - ) - - assert ( - completion_stream is not None - ), f"Failed to handle Unicode character: {unicode_char}" - - -if __name__ == "__main__": - test_utf8_streaming_sync() - asyncio.run(test_utf8_streaming_async()) - test_various_unicode_characters() - print("All UTF-8 streaming tests passed!") + assert "".join([chunk["text"] async for chunk in completion_stream]) == text diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index f78c6064f76..3336cce9aaf 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -74,6 +74,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/proxy/container_endpoints/ownership.py", "litellm/proxy/discovery_endpoints/agent_skills_endpoints.py", "litellm/proxy/guardrails/guardrail_hooks/straiker/straiker.py", + "litellm/proxy/hooks/model_budget_shared_windows.py", "litellm/proxy/spend_tracking/key_metadata_recovery.py", "litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py", "litellm/responses/litellm_completion_transformation/transformation.py", diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index 36be864a6f7..195d5afe000 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -862,8 +862,6 @@ def test_secret_filter_keeps_truncated_traceback(monkeypatch): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_redaction_precedes_a_credential_cut(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setenv("MAX_STRING_LENGTH_STDOUT_LOG", "500") monkeypatch.setenv("MAX_BASE64_LENGTH_STDOUT_LOG", "0") @@ -890,8 +888,6 @@ def test_correlation_id_redacts_before_its_length_bound(monkeypatch): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.WARNING, "bad % api_key=secret123", ("value",)) @@ -905,8 +901,6 @@ def test_malformed_interpolation_still_scrubs_a_record(monkeypatch, native): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_key_pattern_template_keeps_the_rendered_redacted_line(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.INFO, "password=%s ok", ("hunter2",)) @@ -1733,8 +1727,6 @@ def test_access_log_path_filter_keeps_a_record_without_a_string_path_arg(monkeyp @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) secret = "sk-" + "q" * 48 @@ -1761,8 +1753,6 @@ def test_diagnostic_filter_scrubs_exc_stack_and_nested_extras(monkeypatch, nativ @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) record = _make_record(logging.WARNING, "api_key=secret123") @@ -1775,8 +1765,6 @@ def test_diagnostic_filter_stamps_records_so_a_second_pass_is_free(monkeypatch, @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) secret = "sk-" + "q" * 48 @@ -1791,8 +1779,6 @@ def test_json_formatter_scrubs_unfiltered_extras(monkeypatch, native): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_filter_redacts_a_non_string_message_object(monkeypatch, native): - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") monkeypatch.setattr("litellm._logging._ENABLE_SECRET_REDACTION", True) secret = "sk-" + "q" * 48 diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index ba39fc3f5ac..f38bbcf4404 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -2127,6 +2127,330 @@ def test_responses_bridge_preserves_reasoning_effort_with_drop_params( assert request_body["reasoning"] == {"effort": "high"} +@pytest.fixture +def openai_responses_route( + restore_model_registry, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> respx.Route: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("LITELLM_REASONING_AUTO_SUMMARY", raising=False) + litellm.register_model( + { + "gpt-6.1-terra": { + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "supports_minimal_reasoning_effort": False, + "reasoning_modes": ["standard", "pro"], + "reasoning_contexts": ["auto", "current_turn", "all_turns"], + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + }, + persist_across_reloads=False, + ) + return respx_mock.post("https://api.openai.com/v1/responses").respond( + json={ + "id": "resp_test", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-6.1-terra", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + } + ) + + +@pytest.mark.parametrize( + ("reasoning_effort", "summary_alias", "expected_reasoning"), + [ + ( + {"effort": "low", "mode": "pro", "context": "current_turn"}, + {}, + {"effort": "low", "mode": "pro", "context": "current_turn"}, + ), + ( + {"effort": "low", "mode": "pro", "context": "current_turn"}, + {"reasoning_summary": "auto"}, + {"effort": "low", "mode": "pro", "context": "current_turn", "summary": "auto"}, + ), + ( + {"effort": "low", "mode": "pro", "summary": "concise"}, + {"reasoningSummary": "auto"}, + {"effort": "low", "mode": "pro", "summary": "concise"}, + ), + ( + {"effort": "low", "context": "current_turn", "summary": None}, + {"reasoning_summary": "auto"}, + {"effort": "low", "context": "current_turn", "summary": None}, + ), + ({"mode": "pro", "context": "all_turns"}, {}, {"mode": "pro", "context": "all_turns"}), + ("low", {}, {"effort": "low"}), + ], +) +def test_openai_responses_bridge_sends_the_whole_reasoning_effort_dict( + reasoning_effort: str | dict[str, str | None], + summary_alias: dict[str, str], + expected_reasoning: dict[str, str | None], + openai_responses_route: respx.Route, +): + caller_reasoning_effort: Final = copy.deepcopy(reasoning_effort) + + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + api_key="fake-key", + **summary_alias, + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == expected_reasoning + assert reasoning_effort == caller_reasoning_effort + + +@pytest.mark.parametrize( + "reasoning_effort, summary_alias, expected_reasoning", + [ + ( + {"effort": "minimal", "mode": "pro", "context": "current_turn"}, + {}, + {"mode": "pro", "context": "current_turn"}, + ), + ({"effort": "minimal", "mode": "pro", "summary": "auto"}, {}, {"mode": "pro", "summary": "auto"}), + ({"effort": "minimal", "mode": "pro"}, {"reasoning_summary": "auto"}, {"mode": "pro", "summary": "auto"}), + ], +) +def test_openai_responses_bridge_drops_only_a_refused_effort_under_drop_params( + reasoning_effort: dict[str, str], + summary_alias: dict[str, str], + expected_reasoning: dict[str, str], + openai_responses_route: respx.Route, +): + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + **summary_alias, + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == expected_reasoning + + +def test_openai_responses_bridge_drops_a_reasoning_effort_dict_for_a_model_without_reasoning( + openai_responses_route: respx.Route, +): + litellm.completion( + model="openai/responses/gpt-4.1", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "current_turn"}, + drop_params=True, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert "reasoning" not in request_body + + +@pytest.mark.parametrize( + "drop_params, expected_reasoning", + [ + (True, {"effort": "low", "context": "current_turn"}), + ( + False, + { + "effort": "low", + "exclude": False, + "mode": "turbo", + "context": "current_turn", + "generate_summary": "concise", + }, + ), + ], +) +def test_openai_responses_bridge_drops_reasoning_fields_openai_refuses_only_under_drop_params( + drop_params: bool, + expected_reasoning: dict[str, object], + openai_responses_route: respx.Route, +): + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={ + "effort": "low", + "exclude": False, + "mode": "turbo", + "context": "current_turn", + "generate_summary": "concise", + }, + drop_params=drop_params, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == expected_reasoning + + +@pytest.mark.parametrize( + "declared, drop_params, expected_reasoning", + [ + ( + {"reasoning_modes": ["standard"], "reasoning_contexts": ["auto", "current_turn"]}, + True, + {"effort": "low"}, + ), + ( + {"reasoning_modes": ["pro"], "reasoning_contexts": ["all_turns"]}, + True, + {"effort": "low", "mode": "pro", "context": "all_turns"}, + ), + ({}, True, {"effort": "low", "mode": "pro", "context": "all_turns"}), + ({"reasoning_modes": ["standard"]}, True, {"effort": "low", "context": "all_turns"}), + ( + {"reasoning_modes": ["standard"], "reasoning_contexts": ["auto", "current_turn"]}, + False, + {"effort": "low", "mode": "pro", "context": "all_turns"}, + ), + ], +) +def test_openai_responses_bridge_drops_a_mode_or_context_the_model_does_not_list_only_under_drop_params( + declared: dict[str, list[str]], + drop_params: bool, + expected_reasoning: dict[str, str], + openai_responses_route: respx.Route, +): + litellm.register_model( + { + "gpt-6.1-reasoning-test": { + "litellm_provider": "openai", + "mode": "chat", + "supports_reasoning": True, + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + **declared, + } + }, + persist_across_reloads=False, + ) + + litellm.completion( + model="openai/responses/gpt-6.1-reasoning-test", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns"}, + drop_params=drop_params, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == expected_reasoning + + +def test_openai_responses_bridge_forwards_mode_and_context_under_drop_params_for_a_catalog_model_with_no_cost_entry( + openai_responses_route: respx.Route, +): + litellm.open_ai_chat_completion_models.add("unpriced-reasoning-model") + + litellm.completion( + model="openai/responses/unpriced-reasoning-model", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "all_turns", "generate_summary": "auto"}, + drop_params=True, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == {"effort": "low", "mode": "pro", "context": "all_turns"} + + +def test_openai_responses_bridge_forwards_a_reasoning_effort_dict_whole_for_a_model_outside_the_openai_catalog( + openai_responses_route: respx.Route, +): + reasoning_effort: Final = {"effort": "high", "exclude": True, "budget_tokens": 2000} + + litellm.completion( + model="openai/responses/my-proxy-alias", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort=reasoning_effort, + drop_params=True, + api_key="fake-key", + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body["reasoning"] == reasoning_effort + + +@pytest.mark.parametrize( + "additional_drop_params, summary_alias, expected_reasoning", + [ + (["reasoning_effort.mode"], {}, {"effort": "low", "context": "current_turn"}), + ( + ["reasoning_effort.summary"], + {"reasoning_summary": "auto"}, + {"effort": "low", "mode": "pro", "context": "current_turn"}, + ), + (["reasoning_effort"], {"reasoning_summary": "auto"}, None), + ("reasoning_effort", {}, {"effort": "low", "mode": "pro", "context": "current_turn"}), + ], +) +def test_openai_responses_bridge_applies_additional_drop_params_to_a_reasoning_effort_dict( + additional_drop_params: list[str] | str, + summary_alias: dict[str, str], + expected_reasoning: dict[str, str] | None, + openai_responses_route: respx.Route, +): + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "current_turn"}, + additional_drop_params=additional_drop_params, + api_key="fake-key", + **summary_alias, + ) + + request_body: Final = json.loads(openai_responses_route.calls[0].request.content) + assert request_body.get("reasoning") == expected_reasoning + + +def test_openai_responses_bridge_gives_a_reasoning_effort_dict_the_auto_summary_a_string_gets( + openai_responses_route: respx.Route, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort="low", + api_key="fake-key", + ) + litellm.completion( + model="openai/responses/gpt-6.1-terra", + messages=[{"role": "user", "content": "hello"}], + reasoning_effort={"effort": "low", "mode": "pro", "context": "current_turn"}, + api_key="fake-key", + ) + + string_reasoning, dict_reasoning = ( + json.loads(call.request.content)["reasoning"] for call in openai_responses_route.calls + ) + assert string_reasoning["summary"] is not None + assert dict_reasoning == { + "effort": "low", + "mode": "pro", + "context": "current_turn", + "summary": string_reasoning["summary"], + } + + _BRIDGED_RESPONSE_ANSWER: Final = json.dumps( { "id": "resp_slow", diff --git a/tests/unit/test_model_prices_and_context_window.py b/tests/unit/test_model_prices_and_context_window.py index bbfe304429d..2aceed493a7 100644 --- a/tests/unit/test_model_prices_and_context_window.py +++ b/tests/unit/test_model_prices_and_context_window.py @@ -115,6 +115,23 @@ def test_github_copilot_backup_rows_match_the_main_catalog() -> None: assert _github_copilot_catalog_rows(PRICES_PATH) == _github_copilot_catalog_rows(BACKUP_PRICES_PATH) +class _TokenLimits(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + max_input_tokens: int + max_output_tokens: int + + +@pytest.mark.parametrize("path", (PRICES_PATH, BACKUP_PRICES_PATH), ids=("main", "backup")) +def test_zerogpu_rows_have_an_output_limit_that_fits_the_context_window(path: Path) -> None: + prices: Final[dict[str, JsonValue]] = TypeAdapter(dict[str, JsonValue]).validate_json(path.read_bytes()) + rows: Final = MappingProxyType( + {key: _TokenLimits.model_validate(value) for key, value in prices.items() if key.startswith("zerogpu/")} + ) + assert rows, f"No ZeroGPU rows found in {path}" + assert tuple(key for key, row in rows.items() if row.max_output_tokens > row.max_input_tokens) == () + + def test_github_copilot_rows_resolve_through_get_model_info() -> None: rows: Final[Mapping[str, _CopilotEndpointRow]] = _copilot_endpoint_rows(PRICES_PATH) assert { diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index 1d9c2e87100..d0c5c8247d3 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -1969,7 +1969,7 @@ ALIAS_CASES: Final = ( ("litellm.router_utils.add_retry_fallback_headers", "HiddenParamsAsyncIteratorWrapper", "_inner", "inner", True), ("litellm.router_utils.add_retry_fallback_headers", "", "_HiddenParamsHost", "HiddenParamsHost", False), ( - "litellm.router_utils.batch_utils", + "litellm.litellm_core_utils.batch_utils", "", "_get_router_metadata_variable_name", "get_router_metadata_variable_name", diff --git a/tests/unit/test_repo_agent_skills.py b/tests/unit/test_repo_agent_skills.py new file mode 100644 index 00000000000..591241f0cf9 --- /dev/null +++ b/tests/unit/test_repo_agent_skills.py @@ -0,0 +1,68 @@ +from collections.abc import Iterator +from fnmatch import fnmatch +from pathlib import Path +from typing import Final + +import pytest +import yaml + +_REPO_ROOT: Final = Path(__file__).resolve().parents[2] +_SKILLS_DIR: Final = _REPO_ROOT / "skills" +_SKILL_DIRECTORIES: Final = tuple(sorted(skill_file.parent for skill_file in _SKILLS_DIR.glob("*/SKILL.md"))) +_FULL_DEPTH_MAX_DEPTH: Final = 5 +_FULL_DEPTH_SKIPPED_NAMES: Final = frozenset(("node_modules", ".git", "dist", "build", "__pycache__")) +_GITIGNORED_NAME_PATTERNS: Final = tuple( + line.strip().rstrip("/") + for line in (_REPO_ROOT / ".gitignore").read_text(encoding="utf-8").splitlines() + if line.strip() and not line.startswith(("#", "!")) and "/" not in line.strip().rstrip("/") +) + + +def _absent_from_a_clone_or_skipped_by_the_cli(name: str) -> bool: + return name in _FULL_DEPTH_SKIPPED_NAMES or any(fnmatch(name, pattern) for pattern in _GITIGNORED_NAME_PATTERNS) + + +def _full_depth_skill_directories(directory: Path, depth: int) -> Iterator[Path]: + if (directory / "SKILL.md").is_file(): + yield directory + if depth == _FULL_DEPTH_MAX_DEPTH: + return + for child in sorted(directory.iterdir()): + if child.is_dir() and not _absent_from_a_clone_or_skipped_by_the_cli(child.name): + yield from _full_depth_skill_directories(child, depth + 1) + + +_NON_END_USER_SKILL_DIRECTORIES: Final = tuple( + directory for directory in _full_depth_skill_directories(_REPO_ROOT, 0) if directory.parent != _SKILLS_DIR +) + + +def _frontmatter(skill_directory: Path) -> object: + opening, frontmatter, _body = (skill_directory / "SKILL.md").read_text(encoding="utf-8").split("---\n", 2) + assert opening == "", f"{skill_directory.name}/SKILL.md must open with a --- frontmatter block" + return yaml.safe_load(frontmatter) + + +@pytest.mark.parametrize("skill_directory", _SKILL_DIRECTORIES, ids=lambda path: path.name) +def test_skill_frontmatter_has_string_name_matching_its_directory_and_a_description(skill_directory: Path) -> None: + frontmatter: Final = _frontmatter(skill_directory) + assert isinstance(frontmatter, dict), frontmatter + assert frontmatter.get("name") == skill_directory.name, frontmatter + description: Final = frontmatter.get("description") + assert isinstance(description, str) and description.strip(), frontmatter + + +@pytest.mark.parametrize( + "skill_directory", + _NON_END_USER_SKILL_DIRECTORIES, + ids=lambda path: str(path.relative_to(_REPO_ROOT)), +) +def test_skill_outside_skills_folder_is_hidden_from_end_user_skill_installs(skill_directory: Path) -> None: + frontmatter: Final = _frontmatter(skill_directory) + assert isinstance(frontmatter, dict), frontmatter + metadata: Final = frontmatter.get("metadata") + assert isinstance(metadata, dict) and metadata.get("internal") is True, ( + f"{skill_directory} needs `metadata: {{internal: true}}`, or `npx skills add BerriAI/litellm` installs it " + "for end users, from a root dot-folder always and from anywhere with --full-depth (skills CLI 1.7.4, " + "checked 2026-10-28)" + ) diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index 7577064b7f9..2a260fdf0b8 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -6,7 +6,7 @@ regardless of the routing strategy being used. """ import asyncio -from datetime import timedelta +from datetime import datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock import pytest @@ -222,11 +222,9 @@ class TestModelRateLimitingCheck: check.log_success_event(kwargs, None, None, None) - # Verify increment_cache was called - mock_cache.increment_cache.assert_called_once() - _, kwarg_params = mock_cache.increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 + keys = tuple(call.kwargs["key"] for call in mock_cache.increment_cache.call_args_list) + assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys + assert all(call.kwargs["value"] == 50 for call in mock_cache.increment_cache.call_args_list) class TestModelRateLimitingCheckAsync: @@ -253,7 +251,7 @@ class TestModelRateLimitingCheckAsync: """Test that RateLimitError is raised when RPM limit is exceeded (async).""" mock_cache = MagicMock() mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=11) # Over limit + mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[11.0]) # Over limit check = ModelRateLimitingCheck(dual_cache=mock_cache) @@ -274,7 +272,7 @@ class TestModelRateLimitingCheckAsync: """Test that requests are allowed when under the limit (async).""" mock_cache = MagicMock() mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache = AsyncMock(return_value=6) + mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[6.0]) check = ModelRateLimitingCheck(dual_cache=mock_cache) @@ -339,7 +337,7 @@ class TestModelRateLimitingCheckAsync: ): redis_cache = MagicMock() redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) - redis_cache.async_increment = AsyncMock(return_value=2) + redis_cache.async_increment_pipeline = AsyncMock(return_value=[2.0]) check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) with pytest.raises(litellm.RateLimitError) as exc_info: @@ -375,11 +373,9 @@ class TestModelRateLimitingCheckAsync: await check.async_log_success_event(kwargs, None, None, None) - # Verify async_increment_cache was called - mock_cache.async_increment_cache.assert_called_once() - _, kwarg_params = mock_cache.async_increment_cache.call_args - assert "test-id:gpt-4:tpm:" in kwarg_params["key"] - assert kwarg_params["value"] == 50 + keys = tuple(call.kwargs["key"] for call in mock_cache.async_increment_cache.call_args_list) + assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys + assert all(call.kwargs["value"] == 50 for call in mock_cache.async_increment_cache.call_args_list) class TestRouterWithEnforceModelRateLimits: @@ -466,3 +462,265 @@ class TestModelRateLimitConcurrency: assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}" assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}" + + +DAY_ONE = datetime(2026, 10, 28, 23, 59, 30, tzinfo=timezone.utc) +RPD_DEPLOYMENT = { + "litellm_params": {"model": "gpt-4.1-mini", "rpd": 2}, + "model_info": {"id": "daily-test-id"}, + "model_name": "daily-model", +} +RPM_AND_RPD_DEPLOYMENT = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpm": 1, "rpd": 2}} + + +class Clock: + def __init__(self, now: datetime): + self.now = now + + def __call__(self) -> datetime: + return self.now + + +def _pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: + return check.pre_call_check(deployment) + + +async def _async_pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: + return await check.async_pre_call_check(deployment) + + +async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: + await check.async_log_success_event(kwargs, None, None, None) + + +class TestDailyLimits: + def test_pre_call_check_rejects_the_request_over_rpd(self): + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError) as exc_info: + _pre_call(check, RPD_DEPLOYMENT) + + assert "RPD limit=2, current usage=3" in str(exc_info.value) + assert "rpd limit=2. current usage=3. id=daily-test-id, model_group=daily-model" in exc_info.value.response.text + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejects_the_request_over_rpd_and_lets_the_router_retry(self): + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + assert await _async_pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + assert await _async_pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError) as exc_info: + await _async_pre_call(check, RPD_DEPLOYMENT) + + assert "RPD limit=2" in str(exc_info.value) + assert exc_info.value.num_retries is None + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejects_once_the_day_tokens_reach_tpd(self): + deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100}} + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + kwargs = { + "standard_logging_object": { + "model_id": "daily-test-id", + "total_tokens": 60, + "hidden_params": {"litellm_model_name": "gpt-4.1-mini"}, + } + } + + await _async_log_success(check, kwargs) + assert await _async_pre_call(check, deployment) == deployment + await _async_log_success(check, kwargs) + with pytest.raises(litellm.RateLimitError) as exc_info: + await _async_pre_call(check, deployment) + + assert "TPD limit=100, current usage=120" in str(exc_info.value) + + def test_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + + assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + clock.now = DAY_ONE - timedelta(minutes=1) + + assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + clock.now = DAY_ONE - timedelta(minutes=2) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + @pytest.mark.asyncio + async def test_async_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + clock.now = DAY_ONE - timedelta(minutes=1) + + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + clock.now = DAY_ONE - timedelta(minutes=2) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + def test_pre_call_check_rejected_requests_hand_back_rpm_and_rpd(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + + assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + clock.now = DAY_ONE - timedelta(minutes=1) + assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + + clock.now = DAY_ONE - timedelta(minutes=2) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + _pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + @pytest.mark.asyncio + async def test_async_pre_call_check_rejected_requests_hand_back_rpm_and_rpd(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + clock.now = DAY_ONE - timedelta(minutes=1) + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + + clock.now = DAY_ONE - timedelta(minutes=2) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + @pytest.mark.asyncio + async def test_async_pre_call_check_refunds_rpm_and_rpd_in_one_cache_trip(self): + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache_pipeline = AsyncMock(side_effect=[[1.0, 3.0], [0.0, 2.0]]) + check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) + + with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): + await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) + + assert mock_cache.async_increment_cache_pipeline.await_count == 2 + refund = mock_cache.async_increment_cache_pipeline.call_args_list[1].kwargs["increment_list"] + assert [(op["key"], op["increment_value"], op["ttl"]) for op in refund] == [ + ("daily-test-id:gpt-4.1-mini:rpm:23-59", -1, 60), + ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", -1, 86400), + ] + + @pytest.mark.asyncio + async def test_async_pre_call_check_increments_rpm_and_rpd_in_one_cache_trip(self): + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[1.0, 1.0]) + check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) + + assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT + + assert mock_cache.async_increment_cache_pipeline.await_count == 1 + operations = mock_cache.async_increment_cache_pipeline.call_args.kwargs["increment_list"] + assert [(op["key"], op["ttl"]) for op in operations] == [ + ("daily-test-id:gpt-4.1-mini:rpm:23-59", 60), + ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", 86400), + ] + assert not mock_cache.async_increment_cache.called + + def test_pre_call_check_reads_the_shared_tpd_count_from_redis(self): + redis_cache = MagicMock() + redis_cache.get_cache.return_value = 100 + deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100}} + check = ModelRateLimitingCheck(dual_cache=DualCache(redis_cache=redis_cache), clock=Clock(DAY_ONE)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + _pre_call(check, deployment) + + assert "TPD limit=100, current usage=100" in str(exc_info.value) + assert redis_cache.get_cache.call_args.kwargs["key"] == "daily-test-id:gpt-4.1-mini:tpd:2026-10-28" + + def test_rpd_count_starts_fresh_after_utc_midnight(self): + clock = Clock(DAY_ONE) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + _pre_call(check, RPD_DEPLOYMENT) + _pre_call(check, RPD_DEPLOYMENT) + with pytest.raises(litellm.RateLimitError): + _pre_call(check, RPD_DEPLOYMENT) + + clock.now = DAY_ONE + timedelta(seconds=30) + + assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT + with pytest.raises(litellm.RateLimitError) as exc_info: + _pre_call(check, RPD_DEPLOYMENT) + assert "current usage=3" in str(exc_info.value) + + @pytest.mark.parametrize( + ("now", "seconds_to_midnight"), + [ + (DAY_ONE, 30), + (datetime(2026, 10, 28, 0, 0, 0, tzinfo=timezone.utc), 86400), + (datetime(2026, 10, 28, 12, 30, 15, 500000, tzinfo=timezone.utc), 41384), + ], + ) + def test_daily_429_retry_after_is_seconds_until_utc_midnight(self, now: datetime, seconds_to_midnight: int): + deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpd": 0}} + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(now)) + + with pytest.raises(litellm.RateLimitError) as exc_info: + _pre_call(check, deployment) + + assert exc_info.value.response.headers["retry-after"] == str(seconds_to_midnight) + + @pytest.mark.asyncio + async def test_router_serves_the_request_from_the_other_deployment_once_rpd_is_spent(self): + router = Router( + model_list=[ + { + "model_name": "daily-group", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_key": "k", + "mock_response": "a", + "rpd": 1, + "order": 1, + }, + "model_info": {"id": "spent-first"}, + }, + { + "model_name": "daily-group", + "litellm_params": { + "model": "openai/gpt-4.1-mini", + "api_key": "k", + "mock_response": "b", + "order": 2, + }, + "model_info": {"id": "always-free"}, + }, + ], + num_retries=1, + ) + check = ModelRateLimitingCheck(dual_cache=router.cache, clock=Clock(DAY_ONE)) + router.optional_callbacks = [check] + litellm.logging_callback_manager.add_litellm_callback(check) + try: + served = [ + (await router.acompletion(model="daily-group", messages=[{"role": "user", "content": "hi"}])) + .choices[0] + .message.content + for _ in range(4) + ] + finally: + router.discard() + + assert served == ["a", "b", "b", "b"] diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index a5a68271111..0ae22534278 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -859,7 +859,7 @@ class TestModelRateLimitingCheckIOTokens: with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) - assert await dual_cache.async_get_cache(key=rpm_key) == 6 + assert await dual_cache.async_get_cache(key=rpm_key) == 5, "the rejected request hands its rpm count back" assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] assert any("both limit types are enforced" in record.message for record in caplog.records) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 63ee1c43edb..300c9cff6ff 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1126,6 +1126,36 @@ def test_arouter_test_team_model(): assert result is not None +def test_team_public_model_names_of_lists_only_that_teams_live_public_names(): + def team_deployment(deployment_id: str, team_id: str, public_model_name: str) -> DeploymentTypedDict: + return { + "model_name": f"model_name_{team_id}_{deployment_id}", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": {"id": deployment_id, "team_id": team_id, "team_public_model_name": public_model_name}, + } + + router: Final = litellm.Router( + model_list=[ + team_deployment("team-a-1", "team-a", "shared-model"), + team_deployment("team-a-2", "team-a", "shared-model"), + team_deployment("team-a-solo", "team-a", "solo-model"), + team_deployment("team-b-1", "team-b", "team-b-model"), + { + "model_name": "plain-model", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}, + "model_info": {"id": "plain-1"}, + }, + ], + ) + + assert router.team_public_model_names_of("team-a") == frozenset({"shared-model", "solo-model"}) + assert router.team_public_model_names_of("team-c") == frozenset() + + router.delete_deployment(id="team-a-solo") + + assert router.team_public_model_names_of("team-a") == frozenset({"shared-model"}) + + def test_team_model_has_alternatives(): def team_deployment( deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False @@ -3774,6 +3804,132 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t ) +def _responses_sse(*events: Mapping[str, object]) -> str: + return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + + +def _responses_envelope(response_id: str, model: str, status: str, **fields: object) -> dict[str, object]: + return { + "id": response_id, + "object": "response", + "created_at": 1700000000, + "status": status, + "model": model, + "output": [], + **fields, + } + + +@pytest.mark.asyncio +async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fallback_usage( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _SpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + respx_mock.post("https://primary.example/v1/responses").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=_responses_sse( + {"type": "response.created", "response": _responses_envelope("resp_primary", "gpt-4o", "in_progress")}, + *( + { + "type": "response.output_text.delta", + "item_id": "msg_primary", + "output_index": 0, + "content_index": 0, + "delta": word, + } + for word in ("Once ", "upon ", "a ", "time ", "there") + ), + { + "type": "response.failed", + "response": _responses_envelope( + "resp_primary", + "gpt-4o", + "failed", + error={"code": "server_error", "message": "The model had an internal error"}, + ), + }, + ), + ) + ) + respx_mock.post("https://backup.example/v1/responses").mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=_responses_sse( + {"type": "response.created", "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress")}, + { + "type": "response.completed", + "response": _responses_envelope( + "resp_backup", + "gpt-4o-mini", + "completed", + output=[ + { + "type": "message", + "id": "msg_backup", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Once upon a time", "annotations": []}], + } + ], + usage={"input_tokens": 1000, "output_tokens": 700, "total_tokens": 1700}, + ), + }, + ), + ) + ) + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k1", "api_base": "https://primary.example/v1"}, + }, + { + "model_name": "backup", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "k2", + "api_base": "https://backup.example/v1", + }, + }, + ], + fallbacks=[{"primary": ["backup"]}], + num_retries=0, + ) + shared_logging, _ = litellm.utils.function_setup( + original_function="aresponses", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model="primary", + input="Tell me a story", + stream=True, + litellm_call_id="shared-call", + ) + + response: Final = await router.aresponses( + model="primary", + input="Tell me a story", + stream=True, + litellm_call_id="shared-call", + litellm_logging_obj=shared_logging, + ) + _ = [event async for event in response] + await asyncio.wait_for(recorder.succeeded.wait(), timeout=10) + + success: Final = recorder.successes[-1] + assert (success["prompt_tokens"], success["completion_tokens"]) == (1000, 700) + spend_row: Final = recorder.spend_rows[-1] + assert (spend_row["prompt_tokens"], spend_row["completion_tokens"], spend_row["spend"]) == ( + 1000, + 700, + success["response_cost"], + ) + + def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited(): from litellm.types.utils import Usage @@ -20630,6 +20786,37 @@ def test_resolved_litellm_models_answers_through_every_channel_a_request_uses( assert set(_resolution_router().resolved_litellm_models(model_name)) == set(expected) +def test_resolved_litellm_models_base_model_precedence_is_opt_in() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "pooled", + "litellm_params": {"model": "openai/deployment-a", "api_key": "sk-test", "base_model": "gpt-5.5"}, + "model_info": {"base_model": "gpt-6-astra"}, + }, + { + "model_name": "pooled", + "litellm_params": {"model": "openai/deployment-b", "api_key": "sk-test", "base_model": "gpt-5.6-sol"}, + }, + { + "model_name": "team-deployment", + "litellm_params": {"model": "openai/deployment-c", "api_key": "sk-test"}, + "model_info": { + "base_model": "gpt-6-astra", + "team_id": "test-team", + "team_public_model_name": "team-model", + }, + }, + ], + model_group_alias={"nickname": "pooled"}, + ) + assert router.resolved_litellm_models("nickname") == ("openai/deployment-a", "openai/deployment-b") + assert router.resolved_litellm_models("nickname", use_base_model=True) == ("gpt-6-astra", "gpt-5.6-sol") + assert router.resolved_litellm_models("team-model", "test-team") == ("openai/deployment-c",) + assert router.resolved_litellm_models("team-model", "test-team", use_base_model=True) == ("gpt-6-astra",) + assert router.resolved_litellm_models("never-configured", use_base_model=True) == () + + class TestTierParamsTheTargetAccepts: """A tier's litellm_params are applied to every request that tier routes, so one the target cannot take raised UnsupportedParamsError before the request left the proxy, turning the whole @@ -31041,3 +31228,51 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( finally: await GLOBAL_LOGGING_WORKER.flush() await client.close() + + +_FALLBACK_VALIDATION_MODEL_LIST: Final = ( + {"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4", "api_key": "fake-key"}}, + {"model_name": "gpt-3.5", "litellm_params": {"model": "openai/gpt-3.5-turbo", "api_key": "fake-key"}}, +) + + +@pytest.mark.parametrize( + ("router_kwargs", "expected_message"), + [ + ({"fallbacks": [{"gpt-4": "gpt-3.5"}]}, "Fallbacks for 'gpt-4' must be a list of model names, got 'gpt-3.5'"), + ({"default_fallbacks": "gpt-3.5"}, "default_fallbacks must be a list of model names, got 'gpt-3.5'"), + ( + {"content_policy_fallbacks": [{"gpt-4": "gpt-3.5"}]}, + "Fallbacks for 'gpt-4' must be a list of model names, got 'gpt-3.5'", + ), + ], + ids=["string_fallback_target", "string_default_fallbacks", "string_content_policy_fallback_target"], +) +def test_router_rejects_a_string_fallback_target_at_init( + router_kwargs: Mapping[str, object], expected_message: str +) -> None: + with pytest.raises(TypeError, match="must be a list of model names") as excinfo: + Router(model_list=list(_FALLBACK_VALIDATION_MODEL_LIST), **router_kwargs) + + assert expected_message in str(excinfo.value) + + +@pytest.mark.parametrize( + ("router_kwargs", "expected_fallbacks"), + [ + ({"fallbacks": [{"gpt-4": ["gpt-3.5"]}]}, [{"gpt-4": ["gpt-3.5"]}]), + ({"default_fallbacks": ["gpt-3.5"]}, [{"*": ["gpt-3.5"]}]), + ( + {"fallbacks": [{"gpt-4": [{"model": "gpt-3.5", "api_key": "other-key"}]}]}, + [{"gpt-4": [{"model": "gpt-3.5", "api_key": "other-key"}]}], + ), + ({"fallbacks": [{"model": "gpt-3.5"}]}, [{"model": "gpt-3.5"}]), + ], + ids=["list_fallback_target", "list_default_fallbacks", "dict_override_fallback_target", "direct_target_entry"], +) +def test_router_keeps_supported_fallback_shapes( + router_kwargs: Mapping[str, object], expected_fallbacks: list[dict[str, object]] +) -> None: + router: Final = Router(model_list=list(_FALLBACK_VALIDATION_MODEL_LIST), **router_kwargs) + + assert router.fallbacks == expected_fallbacks diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index af8e7f53462..513136ae150 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -4,20 +4,21 @@ import gc import json import os import weakref -from datetime import datetime, timezone from collections.abc import Iterator, Mapping +from datetime import datetime, timezone from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch -import litellm import pytest + +import litellm from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY -from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook from litellm.types.litellm_params import RoutingStrategyName @@ -801,10 +802,11 @@ async def test_deployment_callback_on_failure(model_list): def test_deployment_callback_respects_cooldown_time(model_list): """Ensure per-model cooldown_time is honored even when exception headers are present.""" - import httpx import time from unittest.mock import patch + import httpx + router = Router(model_list=model_list) class FakeException(Exception): @@ -1811,8 +1813,8 @@ def test_get_allowed_fails_from_policy(model_list, exception_type, exception_nam def test_initialize_alerting(model_list): """Test if the 'initialize_alerting' function is working correctly""" - from litellm.types.router import AlertingConfig from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + from litellm.types.router import AlertingConfig router = Router( model_list=model_list, alerting_config=AlertingConfig(webhook_url="test") @@ -2245,6 +2247,7 @@ def test_deployments_by_pattern(model_list): ) def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model): from re import Match + from litellm.router_utils.pattern_match_deployments import PatternMatchRouter pattern_router = PatternMatchRouter() @@ -2702,7 +2705,7 @@ def test_handle_clientside_credential_metadata_loading(model_list, function_name ) def test_handle_clientside_credential_metadata_variable_name(model_list, function_name, metadata_key): """Test that _handle_clientside_credential uses the correct metadata variable name based on function name""" - from litellm.router_utils.batch_utils import get_router_metadata_variable_name + from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name router = Router(model_list=model_list) diff --git a/tests/unit/test_router/test_router_surface_check.py b/tests/unit/test_router/test_router_surface_check.py new file mode 100644 index 00000000000..a97fd029380 --- /dev/null +++ b/tests/unit/test_router/test_router_surface_check.py @@ -0,0 +1,58 @@ +""" +The router surface fence (tests/code_coverage_tests/router_surface_check.py) must see every way +outside code reaches a Router member, or new dependencies slip past the shrink-only allowlist. +""" + +import ast +import importlib.util +import pathlib + +import pytest + +_CHECK = pathlib.Path(__file__).parents[2] / "code_coverage_tests" / "router_surface_check.py" + + +@pytest.fixture(scope="module") +def surface_check(): + spec = importlib.util.spec_from_file_location("router_surface_check", _CHECK) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _uses(surface_check, src: str) -> set: + return surface_check.member_uses(ast.parse(src), surface_check.router_members()) + + +def test_instance_and_dynamic_access_are_uses(surface_check): + assert _uses(surface_check, "llm_router.get_model_list()\ngetattr(router, 'model_list')") == { + "get_model_list", + "model_list", + } + + +def test_class_level_calls_are_uses(surface_check): + assert _uses(surface_check, "litellm.Router.get_valid_args()") == {"get_valid_args"} + assert _uses(surface_check, "from litellm import Router\nRouter.get_valid_args()") == {"get_valid_args"} + + +def test_container_file_endpoints_registered_by_name_are_members(surface_check): + members = surface_check.router_members() + assert {"aupload_container_file", "upload_container_file", "alist_container_files"} <= members + assert _uses(surface_check, "await llm_router.aupload_container_file(file=f)") == {"aupload_container_file"} + + +def test_unlisted_class_level_use_fails_the_fence(surface_check): + found = { + "members": {"outside": {"get_valid_args": {"litellm/new_module.py"}}, "router_family": {}}, + "router_utils_imports": {}, + } + allow = {"members_outside": [], "members_router_family": [], "router_utils_imports_outside": []} + errors = surface_check.compare(allow, found) + assert len(errors) == 1 + assert "'get_valid_args' in litellm/new_module.py" in errors[0] + + +def test_unrelated_receivers_are_not_uses(surface_check): + assert _uses(surface_check, "client.get_model_list()\nRouterConfig.model_list") == set() diff --git a/tests/unit/test_secret_redaction.py b/tests/unit/test_secret_redaction.py index 3b0e9390a61..cce9c6c558d 100644 --- a/tests/unit/test_secret_redaction.py +++ b/tests/unit/test_secret_redaction.py @@ -76,8 +76,6 @@ def test_redact_string_catches_secret_patterns(): @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) def test_diagnostic_redaction_policy_matches_across_backends(monkeypatch: pytest.MonkeyPatch, native: bool) -> None: - if native: - pytest.importorskip("litellm.rust_bridge._native") monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") assert redact_string("GET /v1?api_key=abcdefgh12345&page=2") == "GET /v1?REDACTED&page=2" diff --git a/tests/unit/test_spend_logs_guard.py b/tests/unit/test_spend_logs_guard.py index 9a033045dd0..eb54099cf8b 100644 --- a/tests/unit/test_spend_logs_guard.py +++ b/tests/unit/test_spend_logs_guard.py @@ -1,8 +1,12 @@ +from __future__ import annotations + import importlib.util import sys from pathlib import Path from typing import Final +import pytest + _SCRIPT: Final = Path(__file__).resolve().parents[2] / "ci_cd" / "spend_logs_guard.py" _SPEC: Final = importlib.util.spec_from_file_location("spend_logs_guard", _SCRIPT) assert _SPEC is not None and _SPEC.loader is not None @@ -10,36 +14,145 @@ guard: Final = importlib.util.module_from_spec(_SPEC) sys.modules["spend_logs_guard"] = guard _SPEC.loader.exec_module(guard) +SCHEMA: Final = ( + "generator client {\n" + ' provider = "prisma-client-py"\n' + "}\n" + "\n" + "model LiteLLM_TeamTable {\n" + " team_id String @id\n" + "}\n" + "\n" + "model LiteLLM_SpendLogs {\n" + " request_id String @id\n" + " api_key String\n" + ' @@index([api_key], map: "LiteLLM_SpendLogs_api_key_idx")\n' + "}\n" + "\n" + "model LiteLLM_UserTable {\n" + " user_id String @id\n" + "}\n" +) +_NO_BLOCKS: Final = {} -def _diff(path: str, added: str, start: int = 10) -> str: - return f"diff --git a/{path} b/{path}\n--- a/{path}\n+++ b/{path}\n@@ -{start},0 +{start},1 @@\n+{added}\n" + +def _hunk(old_start: int, old_count: int, new_start: int, new_count: int, *lines: str) -> str: + body: Final = "".join(f"{line}\n" for line in lines) + return f"@@ -{old_start},{old_count} +{new_start},{new_count} @@\n{body}" + + +def _diff(path: str, *hunks: str) -> str: + return f"diff --git a/{path} b/{path}\n--- a/{path}\n+++ b/{path}\n" + "".join(hunks) + + +def _added(path: str, added: str, start: int = 10) -> str: + return _diff(path, _hunk(start, 0, start, 1, f"+{added}")) + + +def _blocks(path: str, base: str = SCHEMA, head: str = SCHEMA) -> dict[str, guard.Block]: + return {path: guard.Block(base=guard.model_block(base), head=guard.model_block(head))} + + +def _with_line(schema: str, after: int, line: str) -> str: + lines: Final = schema.splitlines(keepends=True) + return "".join(lines[:after]) + f"{line}\n" + "".join(lines[after:]) + + +class TestModelBlock: + def test_the_block_runs_from_the_model_line_through_its_closing_brace(self): + assert guard.model_block(SCHEMA) == range(9, 14) + + def test_a_schema_without_the_model_has_an_empty_block(self): + assert guard.model_block(SCHEMA.replace("LiteLLM_SpendLogs", "LiteLLM_ErrorLogs")) == range(0) class TestSpendLogHits: - def test_an_added_schema_line_naming_the_table_is_a_hit(self): + def test_a_field_added_inside_the_block_is_a_hit_without_naming_the_table(self): + head: Final = _with_line(SCHEMA, 9, " gauntlet_probe String?") + diff: Final = _added("schema.prisma", " gauntlet_probe String?") + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == ( + "schema.prisma: +gauntlet_probe String?", + ) + + def test_a_field_removed_from_the_block_is_a_hit(self): + head: Final = SCHEMA.replace(" api_key String\n", "") + diff: Final = _diff("schema.prisma", _hunk(11, 1, 10, 0, "- api_key String")) + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == ("schema.prisma: -api_key String",) + + def test_the_whole_model_appearing_or_disappearing_is_a_hit(self): + without: Final = SCHEMA.replace( + SCHEMA[SCHEMA.index("model LiteLLM_SpendLogs") : SCHEMA.index("model LiteLLM_UserTable")], "" + ) + body: Final = tuple(SCHEMA.splitlines()[8:14]) + added: Final = _diff("schema.prisma", _hunk(8, 0, 9, 6, *(f"+{line}" for line in body))) + removed: Final = _diff("schema.prisma", _hunk(9, 6, 8, 0, *(f"-{line}" for line in body))) + assert len(guard.spend_log_hits(added, _blocks("schema.prisma", base=without))) == 5 + assert len(guard.spend_log_hits(removed, _blocks("schema.prisma", head=without))) == 5 + + def test_a_line_naming_the_table_in_another_model_is_not_a_hit(self): + head: Final = _with_line(SCHEMA, 6, " spend_logs LiteLLM_SpendLogs[]") diff: Final = _diff( - "schema.prisma", ' @@index([api_key, startTime], map: "LiteLLM_SpendLogs_api_key_startTime_idx")', 740 - ) - assert guard.spend_log_hits(diff) == ( - 'schema.prisma: @@index([api_key, startTime], map: "LiteLLM_SpendLogs_api_key_startTime_idx")', + "schema.prisma", + _hunk(6, 0, 7, 1, "+ spend_logs LiteLLM_SpendLogs[]"), + _hunk(16, 1, 17, 0, "- user_id String @id"), ) + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == () - def test_an_added_migration_line_is_a_hit_in_any_case(self): + def test_each_hunk_is_placed_by_its_own_header(self): + head: Final = _with_line(_with_line(SCHEMA, 2, ' output = "./client"'), 10, " new_col String?") + diff: Final = _diff( + "schema.prisma", + _hunk(2, 0, 3, 1, '+ output = "./client"'), + _hunk(9, 0, 11, 1, "+ new_col String?"), + ) + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == ("schema.prisma: +new_col String?",) + + def test_an_added_migration_line_is_a_hit_in_any_case_and_a_removed_one_is_not(self): path: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/20260101000000_x/migration.sql" - hits: Final = guard.spend_log_hits(_diff(path, 'CREATE INDEX "i" ON "litellm_spendlogs"("api_key");', 3)) - assert hits == (f'{path}: CREATE INDEX "i" ON "litellm_spendlogs"("api_key");',) + diff: Final = _diff( + path, + _hunk( + 3, + 1, + 3, + 1, + '-DROP INDEX "LiteLLM_SpendLogs_old_idx";', + '+CREATE INDEX "i" ON "litellm_spendlogs"("api_key");', + ), + ) + assert guard.spend_log_hits(diff, _NO_BLOCKS) == ( + f'{path}: +CREATE INDEX "i" ON "litellm_spendlogs"("api_key");', + ) - def test_a_removed_line_or_another_table_or_an_unwatched_file_is_not_a_hit(self): - removed: Final = "--- a/schema.prisma\n+++ b/schema.prisma\n@@ -1,1 +1,0 @@\n-model LiteLLM_SpendLogs {\n" - other: Final = _diff("schema.prisma", " team_id String // LiteLLM_TeamTable, not MyLiteLLM_SpendLogs", 5) - elsewhere: Final = _diff("litellm/proxy/utils.py", 'table = "LiteLLM_SpendLogs"', 5) - assert guard.spend_log_hits(removed + other + elsewhere) == () + def test_another_table_in_a_migration_or_an_unwatched_file_is_not_a_hit(self): + path: Final = "litellm-proxy-extras/litellm_proxy_extras/migrations/20260101000000_x/migration.sql" + other: Final = _added( + path, 'ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN "x" TEXT; -- not MyLiteLLM_SpendLogs', 5 + ) + elsewhere: Final = _added("litellm/proxy/utils.py", 'table = "LiteLLM_SpendLogs"', 5) + assert guard.spend_log_hits(other + elsewhere, _NO_BLOCKS) == () + + def test_a_table_whose_name_only_starts_with_the_watched_one_is_not_a_hit(self): + index: Final = ' @@index([team_id], map: "LiteLLM_SpendLogs_team_id_idx")' + head: Final = _with_line(_with_line(SCHEMA, 18, "model LiteLLM_SpendLogsArchive {"), 12, index) + diff: Final = _diff( + "schema.prisma", + _hunk(12, 0, 13, 1, f"+{index}"), + _hunk(18, 0, 20, 1, "+model LiteLLM_SpendLogsArchive {"), + ) + assert guard.spend_log_hits(diff, _blocks("schema.prisma", head=head)) == ( + 'schema.prisma: +@@index([team_id], map: "LiteLLM_SpendLogs_team_id_idx")', + ) def test_two_files_in_one_diff_are_attributed_to_their_own_path(self): - diff: Final = _diff("litellm/proxy/schema.prisma", " a LiteLLM_SpendLogs", 1) + _diff( - "schema.prisma", " b String", 1 - ) - assert guard.spend_log_hits(diff) == ("litellm/proxy/schema.prisma: a LiteLLM_SpendLogs",) + head: Final = _with_line(SCHEMA, 9, " a String") + diff: Final = _added("litellm/proxy/schema.prisma", " a String") + _added("schema.prisma", " b String", 2) + blocks: Final = _blocks("litellm/proxy/schema.prisma", head=head) | _blocks("schema.prisma") + assert guard.spend_log_hits(diff, blocks) == ("litellm/proxy/schema.prisma: +a String",) + + +def _review(login: str, state: str, commit_id: str | None = "head") -> guard.Review: + return guard.Review(state=state, user=guard.Reviewer(login=login), commit_id=commit_id) class TestApprovedBy: @@ -55,25 +168,90 @@ class TestApprovedBy: assert guard.approved_by(reviews, "kerry-berri") is True assert guard.approved_by(reviews[::-1], "kerry-berri") is False + @pytest.mark.parametrize( + ("states", "approved"), + ( + (("APPROVED", "COMMENTED"), True), + (("APPROVED", "PENDING"), True), + (("CHANGES_REQUESTED", "COMMENTED"), False), + (("APPROVED", "DISMISSED"), False), + (("COMMENTED",), False), + ), + ) + def test_a_comment_or_pending_review_does_not_change_the_answer(self, states, approved): + reviews: Final = tuple(guard.Review(state=state, user=guard.Reviewer(login="kerry-berri")) for state in states) + assert guard.approved_by(reviews, "kerry-berri") is approved + def test_no_review_from_the_approver_is_not_an_approval(self): other: Final = guard.Review(state="APPROVED", user=guard.Reviewer(login="someone-else")) assert guard.approved_by((other,), "kerry-berri") is False assert guard.approved_by((), "kerry-berri") is False +class TestLatestApproval: + def test_the_approvers_last_approved_review_is_returned_with_its_commit(self): + reviews: Final = ( + _review("kerry-berri", "APPROVED", "old"), + _review("someone-else", "APPROVED", "head"), + _review("kerry-berri", "APPROVED", "new"), + _review("kerry-berri", "COMMENTED", "head"), + ) + assert guard.latest_approval(reviews, "kerry-berri") is reviews[2] + assert guard.latest_approval(reviews[:2], "kerry-berri") is reviews[0] + assert guard.latest_approval(reviews[1:2], "kerry-berri") is None + assert guard.latest_approval((), "kerry-berri") is None + + +class TestApprovalCovers: + _HITS: Final = ("schema.prisma: a LiteLLM_SpendLogs", "schema.prisma: b LiteLLM_SpendLogs") + + def test_an_approval_on_the_head_covers_it_without_reading_the_approved_commit(self): + def unreadable(commit: str) -> None: + raise AssertionError(f"read {commit}") + + assert guard.approval_covers(_review("kerry-berri", "APPROVED", "head"), "head", self._HITS, unreadable) is True + + def test_an_older_approval_covers_the_head_only_when_it_saw_every_spend_log_line(self): + approval: Final = _review("kerry-berri", "APPROVED", "old") + assert guard.approval_covers(approval, "head", self._HITS, lambda commit: self._HITS) is True + assert guard.approval_covers(approval, "head", self._HITS[:1], lambda commit: self._HITS) is True + assert guard.approval_covers(approval, "head", self._HITS, lambda commit: self._HITS[:1]) is False + + def test_an_approval_whose_commit_cannot_be_read_does_not_cover_the_head(self): + gone: Final = _review("kerry-berri", "APPROVED", "gone") + assert guard.approval_covers(gone, "head", self._HITS, lambda commit: None) is False + without_commit: Final = _review("kerry-berri", "APPROVED", None) + assert guard.approval_covers(without_commit, "head", self._HITS, lambda commit: self._HITS) is False + + class TestExitCode: def test_a_pull_request_without_spend_log_lines_passes_without_reviews(self, tmp_path, monkeypatch, capsys): - monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " foo String", 1)) + monkeypatch.setattr(guard, "added_diff", lambda base, head: _added("schema.prisma", " foo String", 6)) + monkeypatch.setattr(guard, "schema_blocks", lambda base, head: _blocks("schema.prisma")) reviews: Final = tmp_path / "r.json" reviews.write_text("[]") assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 0 assert "nothing to review" in capsys.readouterr().out def test_a_spend_log_change_fails_until_the_approver_approves(self, tmp_path, monkeypatch, capsys): - monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " x LiteLLM_SpendLogs", 7)) + head: Final = _with_line(SCHEMA, 9, " x String?") + monkeypatch.setattr(guard, "added_diff", lambda base_sha, head_sha: _added("schema.prisma", " x String?")) + monkeypatch.setattr(guard, "schema_blocks", lambda base_sha, head_sha: _blocks("schema.prisma", head=head)) reviews: Final = tmp_path / "r.json" reviews.write_text('[{"user": {"login": "someone-else"}, "state": "APPROVED"}]') assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 - assert "schema.prisma: x LiteLLM_SpendLogs" in capsys.readouterr().out - reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED"}]') + assert "schema.prisma: +x String?" in capsys.readouterr().out + reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED", "commit_id": "b"}]') assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 0 + + def test_an_approval_before_a_new_spend_log_line_was_pushed_fails(self, tmp_path, monkeypatch, capsys): + head: Final = _with_line(SCHEMA, 9, " x String?") + monkeypatch.setattr(guard, "added_diff", lambda base_sha, head_sha: _added("schema.prisma", " x String?")) + monkeypatch.setattr(guard, "schema_blocks", lambda base_sha, head_sha: _blocks("schema.prisma", head=head)) + monkeypatch.setattr(guard, "reviewed_hits", lambda base_sha, commit: () if commit == "old" else None) + reviews: Final = tmp_path / "r.json" + reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED", "commit_id": "old"}]') + assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 + assert "changed after @kerry-berri approved old" in capsys.readouterr().out + reviews.write_text('[{"user": {"login": "kerry-berri"}, "state": "APPROVED", "commit_id": "gone"}]') + assert guard.main(["--base", "a", "--head", "b", "--approver", "kerry-berri", "--reviews", str(reviews)]) == 1 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 38287b25169..146e2df6413 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1195,6 +1195,8 @@ _COST_MAP_ROW_SCHEMA: Final = { "type": "array", "items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]}, }, + "reasoning_modes": {"type": "array", "items": {"type": "string"}}, + "reasoning_contexts": {"type": "array", "items": {"type": "string"}}, "default_reasoning_effort": { "type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 740342f54d9..cc1ea645012 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -154,6 +154,8 @@ OPTION_NAMES: Final = ( "model_info", "rpm", "tpm", + "rpd", + "tpd", "itpm", "otpm", "default_api_key_rpm_limit", diff --git a/tests/windows_tests/check_windows_wheel_install.py b/tests/windows_tests/check_windows_wheel_install.py index a6c2e7f2984..b554221030d 100644 --- a/tests/windows_tests/check_windows_wheel_install.py +++ b/tests/windows_tests/check_windows_wheel_install.py @@ -87,10 +87,20 @@ def main(argv): f"::error::installing {os.path.basename(wheel)} into a deep prefix failed" ) return 1 - if _run([python, "-c", "import litellm; import litellm.types.utils"]) != 0: + if _run([python, "-I", "-c", "import litellm; import litellm.types.utils"]) != 0: print("::error::litellm did not import after install (half-unpacked package)") return 1 + checker = os.path.abspath(os.path.join("tests", "base_sdk_tests", "check_base_sdk_install.py")) + for rust in ("0", "1"): + result = subprocess.run( + [python, "-I", checker, "--profile", "native"], + env={**os.environ, "LITELLM_RUST": rust, "LITELLM_LOCAL_MODEL_COST_MAP": "True"}, + check=False, + ) + if result.returncode != 0: + return result.returncode + print( f"ok: {os.path.basename(wheel)} installs into a worst-case prefix and imports" ) diff --git a/ui/litellm-dashboard/.agents/skills/url-state/SKILL.md b/ui/litellm-dashboard/.agents/skills/url-state/SKILL.md index 986118d81a7..8bf8fadd19e 100644 --- a/ui/litellm-dashboard/.agents/skills/url-state/SKILL.md +++ b/ui/litellm-dashboard/.agents/skills/url-state/SKILL.md @@ -1,6 +1,8 @@ --- name: url-state description: Read or write dashboard URL query state (tabs, filters, pagination, deep links, demo flags) with nuqs instead of raw history or URLSearchParams +metadata: + internal: true --- # URL state diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index e67c1105d86..9fce93782b2 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -498,9 +498,6 @@ } }, "src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx": { - "no-nested-ternary": { - "count": 1 - }, "react-hooks/set-state-in-effect": { "count": 1 } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx index d19374f48ca..183d68da261 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx @@ -2,6 +2,7 @@ import React, { useState } from "react"; import { Copy, Info } from "lucide-react"; import { toast } from "@/lib/toast"; import { Button } from "@/components/ui/button"; +import { Kbd } from "@/components/ui/kbd"; import { Textarea } from "@/components/ui/textarea"; import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; @@ -165,9 +166,7 @@ export function GuardrailTestPanel({ />
- Press Enter to - submit • Shift+Enter{" "} - for new line + Press Enter to submit • Shift+Enter for new line Characters: {inputText.length}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.test.ts new file mode 100644 index 00000000000..3ef8e88b2c9 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.test.ts @@ -0,0 +1,118 @@ +/* @vitest-environment jsdom */ +import React from "react"; +import { act, renderHook, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { mcpServerToolKeys, useMCPServerTools } from "./useMCPServerTools"; +import * as networking from "@/components/networking"; + +vi.mock("@/components/networking", () => ({ + listMCPTools: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(() => ({ accessToken: "test-token-123" })), +})); + +const wrapper = ({ children }: { children: React.ReactNode }) => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + return React.createElement(QueryClientProvider, { client: queryClient }, children); +}; + +const toolsFor = (serverId: string) => ({ + tools: [{ name: `${serverId}_tool`, description: `from ${serverId}` }], + error: null, +}); + +describe("useMCPServerTools", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("fetches only the enabled servers and leaves the rest idle", async () => { + vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) => toolsFor(serverId)); + + const { result } = renderHook(() => useMCPServerTools(["a", "b"], new Set(["a"])), { wrapper }); + + await waitFor(() => expect(result.current.get("a")?.tools).toEqual([{ name: "a_tool", description: "from a" }])); + expect(result.current.get("b")).toEqual({ tools: [], status: "idle" }); + expect(networking.listMCPTools).toHaveBeenCalledTimes(1); + expect(networking.listMCPTools).toHaveBeenCalledWith("test-token-123", "a"); + }); + + it("isolates a failing server from the others", async () => { + vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) => + serverId === "bad" ? { tools: [], error: "http_500", message: "upstream down" } : toolsFor(serverId), + ); + + const { result } = renderHook(() => useMCPServerTools(["good", "bad"], new Set(["good", "bad"])), { wrapper }); + + await waitFor(() => expect(result.current.get("bad")).toEqual({ tools: [], status: "error" })); + expect(result.current.get("good")).toEqual({ + tools: [{ name: "good_tool", description: "from good" }], + status: "success", + }); + }); + + it("reports a server that is being fetched as loading", () => { + vi.mocked(networking.listMCPTools).mockReturnValue(new Promise(() => {})); + + const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper }); + + expect(result.current.get("a")).toEqual({ tools: [], status: "loading" }); + }); + + it("normalizes a bare array response, tools without descriptions, and drops tools without a name", async () => { + vi.mocked(networking.listMCPTools).mockResolvedValue([ + { tool_name: "legacy" }, + "plain", + { description: "nameless" }, + ]); + + const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper }); + + await waitFor(() => + expect(result.current.get("a")?.tools).toEqual([ + { name: "legacy", description: "" }, + { name: "plain", description: "" }, + ]), + ); + }); + + it("reports an error when a refetch fails for a server that had loaded no tools", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const sharedWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + vi.mocked(networking.listMCPTools) + .mockResolvedValueOnce({ tools: [], error: null }) + .mockResolvedValueOnce({ tools: [], error: "http_502", message: "upstream down" }); + + const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper: sharedWrapper }); + await waitFor(() => expect(result.current.get("a")?.status).toBe("success")); + + await act(() => queryClient.refetchQueries({ queryKey: mcpServerToolKeys.detail("a") })); + + await waitFor(() => expect(result.current.get("a")).toEqual({ tools: [], status: "error" })); + }); + + it("keeps showing the loaded tools when a later refetch fails", async () => { + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + const sharedWrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + vi.mocked(networking.listMCPTools) + .mockResolvedValueOnce(toolsFor("a")) + .mockResolvedValueOnce({ tools: [], error: "http_502", message: "upstream down" }); + + const { result } = renderHook(() => useMCPServerTools(["a"], new Set(["a"])), { wrapper: sharedWrapper }); + await waitFor(() => expect(result.current.get("a")?.status).toBe("success")); + + await act(() => queryClient.refetchQueries({ queryKey: mcpServerToolKeys.detail("a") })); + + await waitFor(() => expect(queryClient.getQueryState(mcpServerToolKeys.detail("a"))?.status).toBe("error")); + expect(networking.listMCPTools).toHaveBeenCalledTimes(2); + expect(result.current.get("a")).toEqual({ + tools: [{ name: "a_tool", description: "from a" }], + status: "success", + }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.ts new file mode 100644 index 00000000000..bc4a6696a6e --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerTools.ts @@ -0,0 +1,70 @@ +import { useQueries, type QueryObserverResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { listMCPTools } from "@/components/networking"; +import useAuthorized from "../useAuthorized"; + +export const mcpServerToolKeys = createQueryKeys("mcpServerTools"); + +export interface MCPServerToolEntry { + name: string; + description: string; +} + +export type MCPServerToolsStatus = "idle" | "loading" | "error" | "success"; + +export interface MCPServerToolsState { + tools: MCPServerToolEntry[]; + status: MCPServerToolsStatus; +} + +type RawTool = string | { name?: string; tool_name?: string; description?: string }; + +const toEntry = (tool: RawTool): MCPServerToolEntry => + typeof tool === "string" + ? { name: tool, description: "" } + : { name: tool.name ?? tool.tool_name ?? "", description: tool.description ?? "" }; + +// listMCPTools never throws: it reports failures as `{ tools: [], error }`. +// Throw here so a broken server surfaces as a query error instead of an empty tool list. +export const fetchServerTools = async (accessToken: string, serverId: string): Promise => { + const result = await listMCPTools(accessToken, serverId); + if (!Array.isArray(result) && result?.error) { + throw new Error(result.message || result.error); + } + const toolList: RawTool[] = Array.isArray(result) ? result : result?.tools ?? []; + return toolList.map(toEntry).filter((tool) => tool.name !== ""); +}; + +const toStatus = ({ + status, + fetchStatus, + data, +}: Pick, "status" | "fetchStatus" | "data">): MCPServerToolsStatus => { + if (status === "error" && data !== undefined && data.length > 0) return "success"; + if (status !== "pending") return status; + return fetchStatus === "idle" ? "idle" : "loading"; +}; + +export const useMCPServerTools = ( + serverIds: readonly string[], + enabledServerIds: ReadonlySet, +): ReadonlyMap => { + const { accessToken } = useAuthorized(); + return useQueries({ + queries: serverIds.map((serverId) => ({ + queryKey: mcpServerToolKeys.detail(serverId), + queryFn: () => fetchServerTools(accessToken!, serverId), + enabled: !!accessToken && enabledServerIds.has(serverId), + staleTime: 5 * 60 * 1000, + retry: false, + refetchOnWindowFocus: false, + })), + combine: (results) => + new Map( + results.map((result, idx): [string, MCPServerToolsState] => [ + serverIds[idx], + { tools: result.data ?? [], status: toStatus(result) }, + ]), + ), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx index 81462c6bf9b..ba9f2bfa839 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx @@ -1,5 +1,5 @@ import React from "react"; -import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; @@ -7,6 +7,7 @@ import { MCPToolsetsTab } from "./MCPToolsetsTab"; import * as networking from "@/components/networking"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { mcpServerToolKeys } from "@/app/(dashboard)/hooks/mcpServers/useMCPServerTools"; import { MCPToolset } from "@/components/mcp_tools/types"; vi.mock("@/components/networking", () => ({ @@ -20,6 +21,7 @@ vi.mock("@/components/networking", () => ({ vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPToolsets", () => ({ useMCPToolsets: vi.fn() })); vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: vi.fn() })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "sk-test" }) })); const setup = () => userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); @@ -197,6 +199,221 @@ describe("MCPToolsetsTab create/edit toolset form", () => { }); }); +describe("MCPToolsetsTab tool search and bulk selection", () => { + const servers = [ + { server_id: "gh", alias: "github", server_name: "GitHub" }, + { server_id: "lin", alias: "linear", server_name: "Linear" }, + { server_id: "broken", alias: "flaky", server_name: "Flaky" }, + ]; + const catalogs: Record = { + gh: [ + { name: "create_issue", description: "Open a GitHub issue" }, + { name: "list_pulls", description: "List pull requests" }, + ], + lin: [ + { name: "create_ticket", description: "Create a Linear issue" }, + { name: "list_projects", description: "List projects" }, + ], + }; + + const renderWithServers = (toolsets: MCPToolset[] = []) => { + vi.mocked(useMCPToolsets).mockReturnValue({ + data: toolsets, + isLoading: false, + } as unknown as ReturnType); + vi.mocked(useMCPServers).mockReturnValue({ data: servers, isSuccess: true } as unknown as ReturnType< + typeof useMCPServers + >); + const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + render( + + + , + ); + return queryClient; + }; + + const editToolset = (tools: MCPToolset["tools"]) => + ({ toolset_id: "ts-1", toolset_name: "triage", description: "", tools }) as MCPToolset; + + const expectedUpdate = (tools: MCPToolset["tools"]) => ({ + toolset_id: "ts-1", + toolset_name: "triage", + description: "", + tools, + }); + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(networking.listMCPTools).mockImplementation(async (_token, serverId) => + serverId in catalogs + ? { tools: catalogs[serverId], error: null } + : { tools: [], error: "http_502", message: "upstream down" }, + ); + vi.mocked(networking.createMCPToolset).mockResolvedValue( + {} as Awaited>, + ); + vi.mocked(networking.updateMCPToolset).mockResolvedValue( + {} as Awaited>, + ); + }); + + it("searches tool names and descriptions across collapsed servers and selects every match", async () => { + const user = setup(); + renderWithServers(); + + await openCreate(user); + fireEvent.change(screen.getByPlaceholderText("e.g. github-linear-tools"), { target: { value: "issues" } }); + await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "issue"); + + expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeInTheDocument(); + expect(screen.getByRole("checkbox", { name: /^create_issue/ })).toBeInTheDocument(); + expect(screen.queryByRole("checkbox", { name: /^list_pulls/ })).not.toBeInTheDocument(); + expect(await screen.findByText("2 matching tools")).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Select all matching" })); + await user.click(screen.getByRole("button", { name: "Create Toolset" })); + + await waitFor(() => { + expect(networking.createMCPToolset).toHaveBeenCalledWith("sk-test", { + toolset_name: "issues", + description: "", + tools: [ + { server_id: "gh", tool_name: "create_issue" }, + { server_id: "lin", tool_name: "create_ticket" }, + ], + }); + }); + }); + + it("deselecting all matches keeps selected tools the search hides", async () => { + const user = setup(); + renderWithServers([ + editToolset([ + { server_id: "gh", tool_name: "list_pulls" }, + { server_id: "gh", tool_name: "create_issue" }, + { server_id: "lin", tool_name: "create_ticket" }, + ]), + ]); + + await openEditFor(user); + await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "create"); + expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeChecked(); + + await user.click(screen.getByRole("button", { name: "Deselect all matching" })); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(networking.updateMCPToolset).toHaveBeenCalledWith( + "sk-test", + expectedUpdate([{ server_id: "gh", tool_name: "list_pulls" }]), + ); + }); + }); + + it("selects all of one server's tools from its header and toggles single rows, fetching only that server", async () => { + const user = setup(); + renderWithServers(); + + await openCreate(user); + fireEvent.change(screen.getByPlaceholderText("e.g. github-linear-tools"), { target: { value: "gh-only" } }); + await user.click(screen.getByRole("button", { name: /^github/ })); + expect(await screen.findByRole("checkbox", { name: /^list_pulls/ })).toBeInTheDocument(); + + await user.click(screen.getByRole("checkbox", { name: "Select all tools on github" })); + expect(screen.getByRole("checkbox", { name: /^list_pulls/ })).toBeChecked(); + // Clicking anywhere on a tool row toggles it exactly once. + await user.click(screen.getByText("List pull requests")); + expect(screen.getByRole("checkbox", { name: /^list_pulls/ })).not.toBeChecked(); + await user.click(screen.getByRole("button", { name: "Create Toolset" })); + + await waitFor(() => { + expect(networking.createMCPToolset).toHaveBeenCalledWith("sk-test", { + toolset_name: "gh-only", + description: "", + tools: [{ server_id: "gh", tool_name: "create_issue" }], + }); + }); + expect(networking.listMCPTools).toHaveBeenCalledTimes(1); + }); + + it("keeps and flags tools whose server is unreachable or gone when saving", async () => { + const user = setup(); + const tools = [ + { server_id: "broken", tool_name: "flaky_tool" }, + { server_id: "deleted", tool_name: "old_tool" }, + { server_id: "gh", tool_name: "create_issue" }, + ]; + renderWithServers([editToolset(tools)]); + + await openEditFor(user); + await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "create"); + + expect(await screen.findByText("Couldn't load tools for this server.")).toBeInTheDocument(); + expect(await screen.findByText("server unreachable")).toBeInTheDocument(); + expect(screen.getByText("server not found")).toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(networking.updateMCPToolset).toHaveBeenCalledWith("sk-test", expectedUpdate(tools)); + }); + }); + + it("stops fetching server tools once the editor is closed after a search", async () => { + const user = setup(); + const queryClient = renderWithServers(); + + await openCreate(user); + await user.type(screen.getByRole("textbox", { name: "Search tools or servers" }), "issue"); + expect(await screen.findByRole("checkbox", { name: /^create_ticket/ })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Cancel" })); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + const callsWhileOpen = vi.mocked(networking.listMCPTools).mock.calls.length; + + await act(() => queryClient.invalidateQueries({ queryKey: mcpServerToolKeys.all })); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(networking.listMCPTools).toHaveBeenCalledTimes(callsWhileOpen); + + await openCreate(user); + expect(screen.getByRole("textbox", { name: "Search tools or servers" })).toHaveValue(""); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(networking.listMCPTools).toHaveBeenCalledTimes(callsWhileOpen); + }); + + it("does not refetch a previously expanded server when the editor is reopened", async () => { + const user = setup(); + const queryClient = renderWithServers(); + + await openCreate(user); + await user.click(screen.getByRole("button", { name: /^github/ })); + expect(await screen.findByRole("checkbox", { name: /^list_pulls/ })).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Cancel" })); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + await act(() => queryClient.invalidateQueries({ queryKey: mcpServerToolKeys.all })); + + await openCreate(user); + await new Promise((resolve) => setTimeout(resolve, 50)); + + expect(networking.listMCPTools).toHaveBeenCalledTimes(1); + expect(screen.queryByRole("checkbox", { name: /^list_pulls/ })).not.toBeInTheDocument(); + }); + + it("clears the whole toolset with Clear all", async () => { + const user = setup(); + renderWithServers([editToolset([{ server_id: "gh", tool_name: "create_issue" }])]); + + await openEditFor(user); + await user.click(screen.getByRole("button", { name: "Clear all" })); + expect(screen.getByText("No tools added yet")).toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Save Changes" })); + + await waitFor(() => { + expect(networking.updateMCPToolset).toHaveBeenCalledWith("sk-test", expectedUpdate([])); + }); + }); +}); + describe("MCPToolsetsTab usage guide", () => { it("shows the toolset URL on the gateway host", () => { renderTab(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx index ff0e8f7b6a9..dc1477b79d3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx @@ -1,28 +1,39 @@ -import React, { useState, useCallback } from "react"; +import React, { useState } from "react"; import { z } from "zod"; import { SortingState } from "@tanstack/react-table"; import { Inbox, Plus, X } from "lucide-react"; import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; import { useQueryClient } from "@tanstack/react-query"; -import { DataTable } from "@/components/shared/DataTable"; +import { useDebouncedValue } from "@tanstack/react-pacer/debouncer"; import { - createMCPToolset, - updateMCPToolset, - deleteMCPToolset, - listMCPTools, - getGatewayBaseUrl, -} from "@/components/networking"; -import { MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types"; + useMCPServerTools, + type MCPServerToolEntry, + type MCPServerToolsState, +} from "@/app/(dashboard)/hooks/mcpServers/useMCPServerTools"; +import { DataTable } from "@/components/shared/DataTable"; +import { createMCPToolset, updateMCPToolset, deleteMCPToolset, getGatewayBaseUrl } from "@/components/networking"; +import { MCPServer, MCPToolset, MCPToolsetTool } from "@/components/mcp_tools/types"; import { FieldGroup } from "@/components/ui/field"; import { FormField } from "@/components/shared/form/FormField"; import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { useZodForm } from "@/lib/forms/useZodForm"; import { toast } from "@/lib/toast"; import { displayToolName, getMCPToolsetTableColumns } from "./MCPToolsetTableColumns"; +import { + deselectTools, + filterServerTools, + isToolSelected, + matchesServerName, + selectTools, + selectionState, + toggleTool, + type SelectionState, +} from "./toolsetSelection"; import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; interface MCPToolsetsTabProps { @@ -37,110 +48,191 @@ const toolsetSchema = z.object({ type ToolsetFormValues = z.infer; -interface MCPToolListProps { +const SEARCH_DEBOUNCE_WAIT_MS = 300; + +const IDLE_SERVER_TOOLS: MCPServerToolsState = { tools: [], status: "idle" }; + +const NO_SERVERS: ReadonlySet = new Set(); + +// The modal stays mounted while closed, so fetch only while it is open; otherwise a finished +// search keeps every server's query live and refetches all of them in the background. +const serversToFetch = ( + open: boolean, + searching: boolean, + serverIds: readonly string[], + expandedServerIds: ReadonlySet, +): ReadonlySet => { + if (!open) return NO_SERVERS; + return searching ? new Set(serverIds) : expandedServerIds; +}; + +const checkboxStateProps = (state: SelectionState) => ({ + checked: state === "all", + indeterminate: state === "some", +}); + +interface ServerView { + server: MCPServer; + serverName: string; + state: MCPServerToolsState; + visibleTools: readonly MCPServerToolEntry[]; +} + +const isPending = (state: MCPServerToolsState) => state.status === "idle" || state.status === "loading"; + +// While searching, a server stays listed until its tools load and turn out not to match. +const isHiddenBySearch = ( + search: string, + serverNames: readonly string[], + state: MCPServerToolsState, + visibleTools: readonly MCPServerToolEntry[], +) => { + if (!search.trim() || state.status !== "success") return false; + return visibleTools.length === 0 && !matchesServerName(search, serverNames); +}; + +const buildServerViews = ( + servers: readonly MCPServer[], + toolsByServer: ReadonlyMap, + search: string, +): ServerView[] => + servers.flatMap((server) => { + const state = toolsByServer.get(server.server_id) ?? IDLE_SERVER_TOOLS; + const names = [server.alias ?? "", server.server_name ?? ""]; + const visibleTools = filterServerTools(search, names, state.tools); + if (isHiddenBySearch(search, names, state, visibleTools)) return []; + return [{ server, serverName: server.alias || server.server_name || server.server_id, state, visibleTools }]; + }); + +const summarizeSearch = (serverViews: readonly ServerView[]) => ({ + matchingTools: serverViews.flatMap(({ server, visibleTools }) => + visibleTools.map((tool) => ({ server_id: server.server_id, tool_name: tool.name })), + ), + pendingServerCount: serverViews.filter(({ state }) => isPending(state)).length, +}); + +// Why a selected tool cannot be confirmed against its server. The tool stays selected and is +// still saved; the note only keeps it from looking lost. +const selectedToolIssue = ( + tool: MCPToolsetTool, + serverNameById: ReadonlyMap, + toolsByServer: ReadonlyMap, +): string | undefined => { + if (!serverNameById.has(tool.server_id)) return "server not found"; + const state = toolsByServer.get(tool.server_id); + if (state?.status === "error") return "server unreachable"; + if (state?.status === "success" && !state.tools.some((t) => t.name === tool.tool_name)) return "tool not found"; + return undefined; +}; + +interface ServerToolGroupProps { serverId: string; serverName: string; - accessToken: string | null; + state: MCPServerToolsState; + visibleTools: readonly MCPServerToolEntry[]; + expanded: boolean; + searchActive: boolean; selectedTools: MCPToolsetTool[]; - onToggle: (tool: MCPToolsetTool) => void; + onToggleExpanded: () => void; + onToggleTool: (tool: MCPToolsetTool) => void; + onSetToolsSelected: (tools: MCPToolsetTool[], selected: boolean) => void; } -interface ToolEntry { - name: string; - description?: string; -} +function ServerToolGroup({ + serverId, + serverName, + state, + visibleTools, + expanded, + searchActive, + selectedTools, + onToggleExpanded, + onToggleTool, + onSetToolsSelected, +}: ServerToolGroupProps) { + const selectedCount = selectedTools.filter((t) => t.server_id === serverId).length; + const visibleRefs = visibleTools.map((tool) => ({ server_id: serverId, tool_name: tool.name })); + const groupState = selectionState(selectedTools, visibleRefs); -function MCPToolList({ serverId, serverName, accessToken, selectedTools, onToggle }: MCPToolListProps) { - const [tools, setTools] = useState([]); - const [loading, setLoading] = useState(false); - const [expanded, setExpanded] = useState(false); - - const selectedSet = new Set(selectedTools.filter((t) => t.server_id === serverId).map((t) => t.tool_name)); - - const fetchTools = useCallback(async () => { - if (!accessToken || tools.length > 0) return; - setLoading(true); - try { - const result = await listMCPTools(accessToken, serverId); - const toolList = Array.isArray(result) ? result : result?.tools ?? []; - setTools(toolList.map((t: any) => ({ name: t.name ?? t.tool_name ?? t, description: t.description ?? "" }))); - } catch { - setTools([]); - } finally { - setLoading(false); + const renderBody = () => { + if (isPending(state)) { + return ( +
+ +
+ ); } - }, [accessToken, serverId, tools.length]); - - const handleToggle = () => { - if (!expanded) fetchTools(); - setExpanded(!expanded); + if (state.status === "error") { + return

Couldn't load tools for this server.

; + } + if (visibleTools.length === 0) { + return ( +

+ {searchActive ? "No matching tools on this server." : "No tools found for this server."} +

+ ); + } + return ( +
+ {visibleRefs.map((ref, idx) => { + const tool = visibleTools[idx]; + const selected = isToolSelected(selectedTools, ref); + return ( + + ); + })} +
+ ); }; return (
- - {expanded && ( -
- {loading ? ( -
- -
- ) : tools.length === 0 ? ( -

No tools found for this server.

- ) : ( -
- {tools.map((tool) => { - const selected = selectedSet.has(tool.name); - return ( - - ); - })} -
- )} -
- )} +
+ onSetToolsSelected(visibleRefs, checked)} + /> + +
+ {expanded &&
{renderBody()}
}
); } @@ -149,21 +241,29 @@ interface CreateToolsetModalProps { open: boolean; onClose: () => void; onSave: (name: string, description: string | undefined, tools: MCPToolsetTool[]) => Promise; - accessToken: string | null; initialToolset?: MCPToolset; } -function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset }: CreateToolsetModalProps) { - const form = useZodForm(toolsetSchema, { - defaultValues: { - toolset_name: initialToolset?.toolset_name || "", - description: initialToolset?.description || "", - }, - }); +const toolsetFormDefaults = (toolset?: MCPToolset): ToolsetFormValues => ({ + toolset_name: toolset?.toolset_name || "", + description: toolset?.description || "", +}); + +function CreateToolsetModal({ open, onClose, onSave, initialToolset }: CreateToolsetModalProps) { + const form = useZodForm(toolsetSchema, { defaultValues: toolsetFormDefaults(initialToolset) }); const [selectedTools, setSelectedTools] = useState(initialToolset?.tools || []); const [saving, setSaving] = useState(false); - const [serverSearch, setServerSearch] = useState(""); - const { data: mcpServers = [] } = useMCPServers(); + const [search, setSearch] = useState(""); + const [debouncedSearch] = useDebouncedValue(search, { wait: SEARCH_DEBOUNCE_WAIT_MS }); + const [expandedServerIds, setExpandedServerIds] = useState>(() => new Set()); + const { data: mcpServers = [], isSuccess: serversLoaded } = useMCPServers(); + const serverIds = React.useMemo(() => mcpServers.map((s) => s.server_id), [mcpServers]); + const searching = search.trim() !== "" && debouncedSearch.trim() !== ""; + const fetchServerIds = React.useMemo( + () => serversToFetch(open, searching, serverIds, expandedServerIds), + [open, searching, serverIds, expandedServerIds], + ); + const toolsByServer = useMCPServerTools(serverIds, fetchServerIds); const serverPrefixById = React.useMemo( () => new Map(mcpServers.map((s) => [s.server_id, s.alias || s.server_name || s.server_id])), [mcpServers], @@ -171,41 +271,54 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset React.useEffect(() => { if (open) { - form.reset({ - toolset_name: initialToolset?.toolset_name || "", - description: initialToolset?.description || "", - }); + form.reset(toolsetFormDefaults(initialToolset)); setSelectedTools(initialToolset?.tools || []); - setServerSearch(""); + setSearch(""); + setExpandedServerIds(new Set()); } }, [open, initialToolset, form]); const handleToggleTool = (tool: MCPToolsetTool) => { - setSelectedTools((prev) => { - const exists = prev.some((t) => t.server_id === tool.server_id && t.tool_name === tool.tool_name); - return exists - ? prev.filter((t) => !(t.server_id === tool.server_id && t.tool_name === tool.tool_name)) - : [...prev, tool]; - }); + setSelectedTools((prev) => toggleTool(prev, tool)); + }; + + const handleSetToolsSelected = (tools: MCPToolsetTool[], selected: boolean) => { + setSelectedTools((prev) => (selected ? selectTools(prev, tools) : deselectTools(prev, tools))); + }; + + const handleToggleExpanded = (serverId: string) => { + setExpandedServerIds((prev) => + prev.has(serverId) ? new Set([...prev].filter((id) => id !== serverId)) : new Set([...prev, serverId]), + ); + }; + + // The create modal stays mounted, so state left from the last session would fetch servers + // on the first render after reopening, before the open effect resets it. + const handleClose = () => { + setSearch(""); + setExpandedServerIds(new Set()); + onClose(); }; const handleSubmit = async (values: ToolsetFormValues) => { setSaving(true); try { await onSave(values.toolset_name, values.description, selectedTools); - onClose(); + handleClose(); } finally { setSaving(false); } }; - const filteredServers = mcpServers.filter((s) => { - const q = serverSearch.toLowerCase(); - return !q || (s.alias || "").toLowerCase().includes(q) || (s.server_name || "").toLowerCase().includes(q); - }); + const searchActive = search.trim() !== ""; + const serverViews = buildServerViews(mcpServers, toolsByServer, search); + const { matchingTools, pendingServerCount } = searchActive + ? summarizeSearch(serverViews) + : { matchingTools: [], pendingServerCount: 0 }; + const matchingState = selectionState(selectedTools, matchingTools); return ( - !open && onClose()}> + !open && handleClose()}> {initialToolset ? "Edit Toolset" : "New Toolset"} @@ -229,32 +342,65 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset setServerSearch(e.target.value)} + placeholder="Search tools or servers..." + aria-label="Search tools or servers" + value={search} + onChange={(e) => setSearch(e.target.value)} /> - {serverSearch && ( + {search && ( - setServerSearch("")}> + setSearch("")}> )} + {searchActive && ( +
+

+ {matchingTools.length} matching {matchingTools.length === 1 ? "tool" : "tools"} + {pendingServerCount > 0 && + ` · loading ${pendingServerCount} ${pendingServerCount === 1 ? "server" : "servers"}…`} +

+
+ + +
+
+ )}
- {filteredServers.length === 0 ? ( + {serverViews.length === 0 ? (

- {mcpServers.length === 0 ? "No MCP servers configured" : "No servers match your search"} + {mcpServers.length === 0 ? "No MCP servers configured" : "No tools or servers match your search"}

) : ( - filteredServers.map((server) => ( - ( + handleToggleExpanded(server.server_id)} + onToggleTool={handleToggleTool} + onSetToolsSelected={handleSetToolsSelected} /> )) )} @@ -266,41 +412,60 @@ function CreateToolsetModal({ open, onClose, onSave, accessToken, initialToolset {/* Right panel: Your Toolset */}
-

- Your Toolset{" "} - ({selectedTools.length} tools) -

+
+

+ Your Toolset{" "} + ({selectedTools.length} tools) +

+ +
{selectedTools.length === 0 ? (

No tools added yet

) : ( - selectedTools.map((tool, idx) => ( -
- - ✕ - - - )) + + ); + }) )}
-
- ↵ Open + ↵ Open - ↑↓ Navigate + ↑↓ Navigate - esc Close + esc Close
diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx index 0a4c3fcda18..6fa362f4509 100644 --- a/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPaletteHint.tsx @@ -2,6 +2,7 @@ import { Search, X } from "lucide-react"; import { z } from "zod"; +import { Kbd } from "@/components/ui/kbd"; import { storageKey } from "@/lib/storage"; import type { PaletteScope } from "./utils"; import { useCommandPalette } from "./CommandPaletteProvider"; @@ -26,9 +27,7 @@ export function CommandPaletteHint({ scope, onDismiss }: CommandPaletteHintProps ); } diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx index 33f09092cca..955b911ca97 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.integration.test.tsx @@ -1,5 +1,6 @@ +import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; -import { renderWithProviders, screen } from "../../../tests/test-utils"; +import { fireEvent, renderWithProviders, screen } from "../../../tests/test-utils"; import { MODEL_MAX_BUDGET_PREMIUM_HINT, ModelMaxBudgetEditor, type ModelMaxBudget } from "./ModelMaxBudgetEditor"; const STORED: ModelMaxBudget = { "gpt-4o": { budget_limit: 5, time_period: "30d" } }; @@ -67,3 +68,65 @@ describe("ModelMaxBudgetEditor with an enterprise license", () => { expect(addButton()).toBeEnabled(); }); }); + +describe("ModelMaxBudgetEditor model families", () => { + it("stores a typed pattern that matches no listed model as one budget entry", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + renderWithProviders( + , + ); + + await user.click(addButton()); + const modelField = screen.getByPlaceholderText("Select model"); + await user.type(modelField, "claude-opus*"); + await user.click(await screen.findByRole("option", { name: /claude-opus\*/ })); + fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "10" } }); + + expect(onChange).toHaveBeenLastCalledWith({ "claude-opus*": { budget_limit: 10, time_period: "30d" } }); + }); + + it("does not offer a typed pattern that another row already budgets", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(addButton()); + const [, newRowModelField] = screen.getAllByPlaceholderText("Select model"); + await user.type(newRowModelField, "claude-opus*"); + + expect(screen.queryByRole("option", { name: /claude-opus\*/ })).not.toBeInTheDocument(); + await user.clear(newRowModelField); + await user.type(newRowModelField, "claude-sonnet*"); + expect(await screen.findByRole("option", { name: /claude-sonnet\*/ })).toBeInTheDocument(); + }); + + it("does not offer a typed name that is neither listed nor a pattern", async () => { + const user = userEvent.setup(); + renderWithProviders( + , + ); + + await user.click(addButton()); + await user.type(screen.getByPlaceholderText("Select model"), "claude-opsu-4-8"); + + expect(screen.queryByRole("option", { name: /claude-opsu-4-8/ })).not.toBeInTheDocument(); + expect(screen.getByText("No models found")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx index a6e72068463..02fc022e881 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx @@ -50,6 +50,8 @@ export const MODEL_BUDGET_PERIOD_OPTIONS = [ const DEFAULT_PERIOD = "30d"; +const MODEL_PATTERN_VALUE = { hint: "Model pattern", accepts: (typed: string) => typed.includes("*") }; + export const entriesToModelMaxBudget = (entries: readonly ModelBudgetEntry[]): ModelMaxBudget => Object.fromEntries( entries @@ -121,7 +123,7 @@ export function ModelMaxBudgetEditor({ const blurb = (
{premiumUser - ? "Cap spend per model over its own window. A budget set on the bare model name also covers the provider-prefixed spelling of that model." + ? "Cap spend per model over its own window. A budget set on the bare model name also covers the provider-prefixed spelling of that model. Pick an access group, or type a pattern like claude-opus*, to share one budget across every model it covers." : MODEL_MAX_BUDGET_PREMIUM_HINT}
); @@ -165,6 +167,8 @@ export function ModelMaxBudgetEditor({ onValueChange={(model) => updateEntry(entry.id, { model })} placeholder="Select model" emptyText="No models found" + customValue={MODEL_PATTERN_VALUE} + reservedValues={entries.flatMap((other) => (other.id !== entry.id && other.model ? [other.model] : []))} disabled={!premiumUser} /> diff --git a/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.test.tsx b/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.test.tsx new file mode 100644 index 00000000000..5c9cf27470e --- /dev/null +++ b/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.test.tsx @@ -0,0 +1,21 @@ +import { render, screen } from "@testing-library/react"; +import { afterEach, describe, expect, it } from "vitest"; +import { setServerRootPath } from "@/lib/serverRootPath"; +import { AriaAvatar } from "./AriaAvatar"; + +describe("AriaAvatar", () => { + afterEach(() => { + setServerRootPath("/"); + }); + + it.each(["/", "/litellm"])("loads the bundled image when the server root is %s", (root) => { + setServerRootPath(root); + + render(); + + expect(screen.getByRole("presentation", { hidden: true })).toHaveAttribute( + "src", + "/_next/static/media/aria-agent.png", + ); + }); +}); diff --git a/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.tsx b/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.tsx index 8c78908d3a9..e0f014624fe 100644 --- a/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.tsx +++ b/ui/litellm-dashboard/src/components/liteadmin/AriaAvatar.tsx @@ -1,10 +1,10 @@ -import { resolveLogoSrc } from "@/lib/assetPaths"; +import ariaAgentImage from "../../../public/assets/aria-agent.png"; import { cn } from "@/lib/cva.config"; export function AriaAvatar({ size = 28, className }: { size?: number; className?: string }) { return ( {ARIA_AGENT_LABEL} - ⌘J + ⌘J ); } diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 0ebd2cdff17..83e80911cff 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -224,6 +224,7 @@ describe("provider_info_helpers", () => { Providers.Sail, Providers.WANDB, Providers.ZAI, + Providers.ZeroGPU, ]; const logolessProviders = Object.values(Providers).filter((provider) => !providerLogoMap[provider]); expect([...logolessProviders].sort()).toEqual([...knownLogolessProviders].sort()); @@ -357,6 +358,10 @@ describe("provider_info_helpers", () => { expect(getPlaceholder(Providers.Cognition)).toBe("cognition/swe-1.7"); }); + it("should return zerogpu/gpt-oss-120b placeholder for ZeroGPU provider", () => { + expect(getPlaceholder("ZeroGPU")).toBe("zerogpu/gpt-oss-120b"); + }); + it("should return a chatgpt/ placeholder for the CHATGPT dropdown key", () => { expect(getPlaceholder("CHATGPT")).toBe("chatgpt/gpt-5.4"); }); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 2541aa45050..bfb578b218a 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -197,6 +197,7 @@ export enum Providers { xAI = "xAI", XINFERENCE = "Xinference", ZAI = "Z.AI (Zhipu AI)", + ZeroGPU = "ZeroGPU", } export const provider_map: Record = { @@ -325,6 +326,7 @@ export const provider_map: Record = { xAI: "xai", XINFERENCE: "xinference", ZAI: "zai", + ZeroGPU: "zerogpu", }; const standaloneSubproviderSlugs = new Set(["bedrock_mantle"]); @@ -497,6 +499,7 @@ const providerPlaceholderMap: Partial> = { [Providers.Voyage]: "voyage/", [Providers.WATSONX]: "watsonx/ibm/granite-3-3-8b-instruct", [Providers.ZAI]: "zai/glm-4.5", + [Providers.ZeroGPU]: "zerogpu/gpt-oss-120b", }; export const getPlaceholder = (selectedProvider: string): string => { diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx index 9f320049b09..ecb8dd161f3 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx @@ -63,6 +63,24 @@ describe("SearchSelect", () => { expect(onValueChange).toHaveBeenLastCalledWith("team-3"); }); + it("offers the typed text as a value only when the caller accepts it", async () => { + const user = userEvent.setup(); + const onValueChange = vi.fn(); + const customValue = { hint: "Custom value", accepts: (typed: string) => typed.endsWith("*") }; + const { rerender } = render(); + const input = screen.getByRole("combobox"); + await user.click(input); + fireEvent.change(input, { target: { value: "team-*" } }); + expect(screen.queryByRole("option", { name: /team-\*/ })).not.toBeInTheDocument(); + + rerender(); + fireEvent.change(input, { target: { value: "team-4" } }); + expect(screen.queryByRole("option", { name: /team-4/ })).not.toBeInTheDocument(); + fireEvent.change(input, { target: { value: "team-*" } }); + await user.click(await screen.findByRole("option", { name: /team-\*/ })); + expect(onValueChange).toHaveBeenLastCalledWith("team-*"); + }); + it("filters the options client-side as you type", async () => { const user = userEvent.setup(); render(); diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx index 3649114826d..e139720a2ed 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.tsx @@ -1,6 +1,6 @@ "use client"; -import React from "react"; +import React, { useState } from "react"; import { Combobox, ComboboxContent, @@ -19,6 +19,11 @@ export interface SearchSelectOption { icon?: React.ReactNode; } +interface CustomValue { + hint: string; + accepts: (typed: string) => boolean; +} + interface SearchSelectProps { options: SearchSelectOption[]; value?: string | null; @@ -29,6 +34,8 @@ interface SearchSelectProps { className?: string; inputId?: string; allowClear?: boolean; + customValue?: CustomValue; + reservedValues?: readonly string[]; "aria-label"?: string; id?: string; } @@ -49,13 +56,22 @@ export function SearchSelect({ className, inputId, allowClear = true, + customValue, + reservedValues = [], "aria-label": ariaLabel, id, }: SearchSelectProps) { + const [query, setQuery] = useState(""); const selected = value == null || value === "" ? null : options.find((option) => option.value === value) ?? { label: value, value }; - const items = + const known = selected !== null && !options.some((option) => option.value === selected.value) ? [selected, ...options] : options; + const typed = query.trim(); + const isNewValue = typed !== "" && !reservedValues.includes(typed) && !known.some((option) => option.value === typed); + const items = + customValue !== undefined && isNewValue && customValue.accepts(typed) + ? [...known, { label: typed, value: typed, sublabel: customValue.hint }] + : known; return ( a.value === b.value} itemToStringLabel={(item: SearchSelectOption) => item.label} filter={matchesQuery} + onInputValueChange={setQuery} disabled={disabled} > ["hotkeys"][number]; @@ -59,14 +60,6 @@ export function shortcutHints(shortcuts: readonly BoundShortcut[]): readonly Sho >((hints, [key, shortcut]) => mergeAdjacent(hints, key, shortcut.description ?? ""), []); } -function Kbd({ children }: { children: ReactNode }) { - return ( - - {children} - - ); -} - /** The shortcuts active right now, read from what is bound rather than typed by hand. */ export function ShortcutHints({ className }: { className?: string }) { const { hotkeys } = useHotkeysContext(); diff --git a/ui/litellm-dashboard/src/components/shared/search/SearchBox.tsx b/ui/litellm-dashboard/src/components/shared/search/SearchBox.tsx index 4e45f8902f0..ae6d5a6ec09 100644 --- a/ui/litellm-dashboard/src/components/shared/search/SearchBox.tsx +++ b/ui/litellm-dashboard/src/components/shared/search/SearchBox.tsx @@ -8,9 +8,10 @@ import { Check, Copy, CornerDownLeft, Loader2, type LucideIcon, Search } from "l import { Schema } from "prosemirror-model"; import { EditorState, Plugin, TextSelection, type Transaction } from "prosemirror-state"; import { Decoration, DecorationSet, type EditorView } from "prosemirror-view"; -import { type ComponentProps, createContext, type ReactNode, useContext, useId, useMemo, useState } from "react"; +import { type ComponentProps, createContext, useContext, useId, useMemo, useState } from "react"; import { Button } from "@/components/ui/button"; +import { Kbd } from "@/components/ui/kbd"; import { cn } from "@/lib/cva.config"; import { type FilterOp, languageOps, parseQuery, type QueryClause, type QueryLanguage } from "./language"; @@ -326,7 +327,7 @@ function Suggestions({ className, children, ...props }: SearchBoxSuggestionsProp - + Select @@ -392,12 +393,4 @@ function OperatorHints({ ops }: { ops: readonly FilterOp[] }) { ); } -function Kbd({ children }: { children: ReactNode }) { - return ( - - {children} - - ); -} - export const SearchBox = { Root, Input, Suggestions, CopyCommand } as const; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f9117afdf3a..a49f1099121 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8475,7 +8475,7 @@ export interface paths { * - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. * - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Supported Claude models on Anthropic, Bedrock, Vertex AI, and Azure AI only. * - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. + * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. * - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. @@ -8721,7 +8721,7 @@ export interface paths { * - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) * - spend: Optional[float] - Amount spent by key * - max_budget: Optional[float] - Max budget for key - * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) * - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. @@ -8798,7 +8798,7 @@ export interface paths { * - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - guardrails: Optional[List[str]] - List of active guardrails for the key * - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. + * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. * - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. @@ -8931,7 +8931,7 @@ export interface paths { * - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) * - spend: Optional[float] - Amount spent by key * - max_budget: Optional[float] - Max budget for key - * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) * - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. Set to null to remove the soft budget. @@ -9020,7 +9020,7 @@ export interface paths { * - tags: Optional[List[str]] - Tags for organizing keys (Enterprise only) * - spend: Optional[float] - Amount spent by key * - max_budget: Optional[float] - Max budget for key - * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} + * - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets keyed by model name, `*` pattern like "claude-opus*", or model access group, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - budget_duration: Optional[str] - Budget reset period ("30d", "1h", etc.) * - soft_budget: Optional[float] - Soft budget limit (warning vs. hard stop). Will trigger a slack alert when this soft budget is reached. @@ -17098,7 +17098,7 @@ export interface paths { * - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. * - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) - * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} + * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} * - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) * - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only. @@ -17328,7 +17328,7 @@ export interface paths { * - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). * - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) - * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} + * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model, keyed by model name, `*` pattern, or model access group. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} * - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) * - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the team. Proxy admin only. @@ -18778,7 +18778,7 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) + * - model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) * - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user. @@ -18917,7 +18917,7 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) + * - model_max_budget: Optional[dict] - Model-specific max budget for user, keyed by model name, `*` pattern, or model access group. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) * - mcp_rpm_limit: Optional[dict] - Per-MCP-server rpm limit, keyed by MCP server name {"github": 100, "slack": 200}. Enforced for keys and teams only; values set on a user are stored but not enforced per user. @@ -23366,9 +23366,10 @@ export interface paths { * -H "Authorization: Bearer sk-..." \ * -H "Content-Type: application/json" \ * -d '{ - * "name": "my-custom-slides-agent", - * "base_agent": "waverunner", - * "instructions": "You are a helpful assistant that creates slides.", + * "id": "my-custom-slides-agent", + * "base_agent": "antigravity-preview-09-2026", + * "system_instruction": "You are a helpful assistant that creates slides.", + * "agent_config": {"type": "antigravity", "model": "gemini-3.8-flash"}, * "base_environment": { * "type": "remote", * "sources": [ @@ -27720,6 +27721,10 @@ export interface components { entities?: { [key: string]: components["schemas"]["MetricWithMetadata"]; }; + /** Guardrails */ + guardrails?: { + [key: string]: components["schemas"]["GuardrailMetric"]; + }; /** Mcp Servers */ mcp_servers?: { [key: string]: components["schemas"]["MetricWithMetadata"]; @@ -33235,6 +33240,38 @@ export interface components { /** Updated At */ updated_at?: string | null; }; + /** GuardrailKeyMetric */ + GuardrailKeyMetric: { + metadata: components["schemas"]["KeyMetadata"]; + metrics: components["schemas"]["GuardrailSpendMetrics"]; + }; + /** GuardrailMetadata */ + GuardrailMetadata: { + /** Guardrail Name */ + guardrail_name: string; + }; + /** GuardrailMetric */ + GuardrailMetric: { + /** Api Key Breakdown */ + api_key_breakdown?: { + [key: string]: components["schemas"]["GuardrailKeyMetric"]; + }; + metadata: components["schemas"]["GuardrailMetadata"]; + metrics: components["schemas"]["GuardrailSpendMetrics"]; + }; + /** GuardrailSpendMetrics */ + GuardrailSpendMetrics: { + /** + * Spend + * @default 0 + */ + spend: number; + /** + * Untracked Units + * @default 0 + */ + untracked_units: number; + }; /** GuardrailSubmissionItem */ GuardrailSubmissionItem: { /** Created At */ @@ -35610,6 +35647,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Responses Api Version */ responses_api_version?: string | null; + /** Rpd */ + rpd?: number | null; /** Rpm */ rpm?: number | null; /** S3 Access Key Id */ @@ -35662,6 +35701,8 @@ export interface components { token_exchange_profile?: string | null; /** Token Exchange Scope */ token_exchange_scope?: string | null; + /** Tpd */ + tpd?: number | null; /** Tpm */ tpm?: number | null; /** Use Chat Completions Api */ @@ -51002,6 +51043,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Responses Api Version */ responses_api_version?: string | null; + /** Rpd */ + rpd?: number | null; /** Rpm */ rpm?: number | null; /** S3 Access Key Id */ @@ -51054,6 +51097,8 @@ export interface components { token_exchange_profile?: string | null; /** Token Exchange Scope */ token_exchange_scope?: string | null; + /** Tpd */ + tpd?: number | null; /** Tpm */ tpm?: number | null; /** Use Chat Completions Api */ @@ -59856,6 +59901,8 @@ export interface operations { purpose?: string | null; limit?: number | null; after?: string | null; + after_id?: string | null; + before_id?: string | null; }; header?: never; path?: never; @@ -79052,6 +79099,8 @@ export interface operations { purpose?: string | null; limit?: number | null; after?: string | null; + after_id?: string | null; + before_id?: string | null; }; header?: never; path?: never; @@ -86387,6 +86436,8 @@ export interface operations { purpose?: string | null; limit?: number | null; after?: string | null; + after_id?: string | null; + before_id?: string | null; }; header?: never; path: {