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/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 355c164a2d4..48407ae532d 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 @@ -1307,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 @@ -1343,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/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..fac9d4ba5d2 100644 --- a/ci_cd/spend_logs_guard.py +++ b/ci_cd/spend_logs_guard.py @@ -4,7 +4,9 @@ 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. +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). The workflow re-runs on every review. """ from __future__ import annotations @@ -13,7 +15,7 @@ import argparse import re import subprocess import sys -from collections.abc import Sequence +from collections.abc import Callable, Sequence from pathlib import Path from typing import Final @@ -25,7 +27,8 @@ 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"(? 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 +111,15 @@ def added_diff(base: str, head: str) -> str: return result.stdout +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)) if fetched(commit) else None + + class Arguments(BaseModel): model_config = ConfigDict(frozen=True) base: str @@ -104,14 +139,21 @@ def main(argv: Sequence[str]) -> int: if not hits: sys.stdout.write("No added line in schema.prisma or a migration names LiteLLM_SpendLogs; 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}The check re-runs when a review is submitted.\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/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index f3b1240f733..47229466021 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -5,9 +5,9 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t import asyncio import json import re -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import replace as dataclasses_replace -from datetime import datetime, timedelta, timezone +from datetime import datetime, timedelta from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast @@ -28,6 +28,7 @@ from litellm.constants import ( from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.repositories.team_repository import TeamRepository from litellm.repositories.verification_token_repository import VerificationTokenRepository +from litellm.utils import get_utc_datetime if TYPE_CHECKING: from prisma import models as prisma_models @@ -42,11 +43,14 @@ if TYPE_CHECKING: def _error_names_resource(error: Exception, resource_id: str) -> bool: - return resource_id in re.findall(r"""[^\s"'`]+""", str(error)) + tokens: Final[list[str]] = re.findall(r"""[^\s"'`<>(]+""", str(error)) + return any(token.rstrip(".,:;)") == resource_id for token in tokens) CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48) +GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) + PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( "completed", @@ -91,6 +95,10 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li return TeamRepository(prisma_client).table +def _unique_jobs(jobs: Sequence[ManagedObjectRow]) -> tuple[ManagedObjectRow, ...]: + return tuple({job.id: job for job in jobs}.values()) + + class CheckBatchCost: def __init__( self, @@ -98,6 +106,7 @@ class CheckBatchCost: prisma_client: "PrismaClient", llm_router: "Router", track_unmanaged_batch_cost: bool = False, + clock: Callable[[], datetime] = get_utc_datetime, ): from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -106,13 +115,14 @@ class CheckBatchCost: self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router self._track_unmanaged_batch_cost = track_unmanaged_batch_cost + self._clock: Final = clock self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckBatchCost") # Cached after the first poll cycle. Once we know the column is absent we skip # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True self.batch_processed_support_confirmed: bool = False self._poll_cursor: tuple[datetime, str] | None = None - self._poll_before: datetime = datetime.now(timezone.utc) + self._poll_before: datetime = clock() self._arrival_cursor: tuple[datetime, str] | None = None self._arrival_before: datetime = self._poll_before self._poll_arrivals: bool = False @@ -147,27 +157,6 @@ class CheckBatchCost: return self.batch_processed_support_confirmed = True - async def _cleanup_stale_managed_objects(self) -> None: - """ - Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days - in non-terminal states as 'stale_expired'. These will never complete and - should not be polled. - """ - cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) - result: Final = await _managed_object_table(self.prisma_client).update_many( - where={ - "file_purpose": "batch", - "status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)}, - "created_at": {"lt": cutoff}, - }, - data={"status": "stale_expired"}, - ) - if result > 0: - verbose_proxy_logger.warning( - f"CheckBatchCost: marked {result} stale managed objects " - f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" - ) - async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]": """Query batch jobs without the batch_processed filter (for older schemas).""" return await self._find_jobs( @@ -187,20 +176,43 @@ class CheckBatchCost: ) async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]: - arrivals: Final = self._poll_arrivals - page: Final = await self._find_job_page(where, arrivals=arrivals) - self._poll_arrivals = not arrivals if page else arrivals - if page: - return page - return await self._find_job_page(where, arrivals=not arrivals) + lead_arrivals: Final = self._poll_arrivals + self._poll_arrivals = not lead_arrivals + lead_share: Final = MAX_OBJECTS_PER_POLL_CYCLE - MAX_OBJECTS_PER_POLL_CYCLE // 2 + lead: Final = await self._find_job_page(where, arrivals=lead_arrivals, take=lead_share) + trail: Final = await self._find_job_page( + where, arrivals=not lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead) + ) + spill: Final = ( + await self._find_job_page( + where, arrivals=lead_arrivals, take=MAX_OBJECTS_PER_POLL_CYCLE - len(lead) - len(trail) + ) + if len(lead) == lead_share + else () + ) + return _unique_jobs((*lead, *trail, *spill)) - async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]: + async def _find_job_page( + self, where: Mapping[str, object], *, arrivals: bool, take: int + ) -> Sequence[ManagedObjectRow]: + if take == 0: + return () + resumed_mid_lane: Final = (self._arrival_cursor if arrivals else self._poll_cursor) is not None + page: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take) + if len(page) == take or not resumed_mid_lane: + return page + wrapped: Final = await self._find_job_page_from_cursor(where, arrivals=arrivals, take=take - len(page)) + return _unique_jobs((*page, *wrapped)) + + async def _find_job_page_from_cursor( + self, where: Mapping[str, object], *, arrivals: bool, take: int + ) -> Sequence[ManagedObjectRow]: cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor if cursor is None: if arrivals: - self._arrival_before = datetime.now(timezone.utc) + self._arrival_before = self._clock() else: - self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW + self._poll_before = self._clock() - RECENT_BATCH_POLL_WINDOW before: Final = self._arrival_before if arrivals else self._poll_before page: Final = await _managed_object_table(self.prisma_client).find_many( where={ @@ -220,10 +232,10 @@ class CheckBatchCost: else {} ), }, - take=MAX_OBJECTS_PER_POLL_CYCLE, + take=take, order=[{"created_at": "asc"}, {"id": "asc"}], ) - next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None + next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == take else None if arrivals: self._arrival_cursor = next_cursor else: @@ -231,11 +243,47 @@ class CheckBatchCost: return page def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]: - return { - "id": job.id, - "updated_at": job.updated_at, - **({"batch_processed": False} if self._has_batch_processed_column else {}), - } + if self._has_batch_processed_column: + return {"id": job.id, "batch_processed": False} + return {"id": job.id, "updated_at": job.updated_at} + + def _is_past_give_up_age(self, job: ManagedObjectRow) -> bool: + return job.created_at < self._clock() - GIVE_UP_POLLING_AFTER + + async def _expire_if_old_and_unfinished(self, job: ManagedObjectRow) -> None: + if job.status in TERMINAL_MANAGED_OBJECT_STATUSES or not self._is_past_give_up_age(job): + return + await self._expire_unfinished(job) + + @staticmethod + def _is_completed_with_output_id_pending(response: "LiteLLMBatch") -> bool: + from litellm.proxy.openai_files_endpoints.common_utils import ( + _completed_batch_safe_to_retire, + ) + + return ( + response.status in ("completed", "complete") + and response.output_file_id is None + and not _completed_batch_safe_to_retire(response) + ) + + async def _expire_unfinished(self, job: ManagedObjectRow) -> None: + try: + expired: Final = await _managed_object_table(self.prisma_client).update_many( + where={ + **self._unprocessed_snapshot(job), + "status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)}, + }, + data={"status": "stale_expired"}, + ) + except Exception as db_err: + verbose_proxy_logger.error(f"CheckBatchCost: failed to mark stale job {job.id} as stale_expired: {db_err}") + return + if expired > 0: + verbose_proxy_logger.warning( + f"CheckBatchCost: job {job.id} is older than {GIVE_UP_POLLING_AFTER.days} days " + "and still not finished, so it is marked stale_expired and no longer polled" + ) async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None: """ @@ -356,15 +404,21 @@ class CheckBatchCost: def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool: """A 404 naming the output file means there is nothing to fetch on this or any later poll: providers like Vertex AI advertise an output path for every batch, - including terminal ones that never wrote it. Any other failure may be - transient, so it keeps retrying until the staleness sweep bounds it.""" + including terminal ones that never wrote it. GCS names a gs:// object without its scheme, + and S3 names an s3:// object by its key alone. Any other failure keeps retrying.""" import openai from litellm.exceptions import NotFoundError + from litellm.llms.base_llm.chat.transformation import BaseLLMException if not output_file_id: return False - return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id) + is_not_found: Final = isinstance(error, (NotFoundError, openai.NotFoundError)) or ( + isinstance(error, BaseLLMException) and error.status_code == 404 + ) + object_path: Final = output_file_id.split("://", 1)[-1] + names: Final = (output_file_id, object_path, object_path.split("/", 1)[-1]) + return is_not_found and any(_error_names_resource(error, name) for name in names) async def _finalize_unbilled_terminal_job( self, job: "ManagedObjectRow", response: "LiteLLMBatch" @@ -1094,13 +1148,6 @@ class CheckBatchCost: processed_models: List[Tuple[Optional[str], Optional[str]]] = [] - try: - await self._cleanup_stale_managed_objects() - except Exception as cleanup_err: - verbose_proxy_logger.warning( - f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}" - ) - # Look for all batches that have not yet been processed by CheckBatchCost. # self._has_batch_processed_column is cached after the first probe so that # older schemas don't pay a guaranteed-failing primary query + warning on @@ -1142,6 +1189,8 @@ class CheckBatchCost: if routing is None: if self._has_unified_id_without_model(job): await self._retire_job(job, "unified object id has no model id") + else: + await self._expire_if_old_and_unfinished(job) continue model_id, batch_id = routing @@ -1166,6 +1215,14 @@ class CheckBatchCost: prom_logger.record_check_batch_cost_error("provider_retrieval_error") if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id): await self._retire_job(job, f"batch {batch_id} no longer exists at the provider") + else: + await self._expire_if_old_and_unfinished(job) + continue + if self._is_completed_with_output_id_pending(response): + verbose_proxy_logger.info( + f"CheckBatchCost: batch {batch_id} is completed but its output file id " + f"has not appeared yet; leaving job {job.id} for the next poll cycle" + ) continue ## RETRIEVE THE BATCH JOB OUTPUT FILE @@ -1217,17 +1274,9 @@ class CheckBatchCost: self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True) elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: - from litellm.proxy.openai_files_endpoints.common_utils import ( - _completed_batch_safe_to_retire, - ) - - if response.status in ("completed", "complete") and not _completed_batch_safe_to_retire(response): - verbose_proxy_logger.info( - f"CheckBatchCost: batch {batch_id} is completed but its output file id " - f"has not appeared yet; leaving job {job.id} for the next poll cycle" - ) - continue await self._retire_unbilled_terminal_job(job, response, model_id, batch_id) + else: + await self._expire_if_old_and_unfinished(job) # Record polling run metrics (always, even if nothing was processed) if prom_logger: diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 48cbc0ab8ff..416444664a3 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -6,6 +6,7 @@ import base64 import json import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from datetime import datetime, timedelta from itertools import chain from types import MappingProxyType from typing import ( @@ -27,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 @@ -39,6 +40,7 @@ from litellm.constants import ( BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS, BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, + MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_FILE_LIST_LIMIT, PENDING_RESPONSE_STATUSES, ) @@ -51,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, @@ -113,6 +115,7 @@ from litellm.types.utils import ( LLMResponseTypes, SpecialEnums, ) +from litellm.utils import get_utc_datetime class _ManagedFileRetrieve(Protocol): @@ -306,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] @@ -359,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: ... @@ -382,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], @@ -507,10 +586,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): prisma_client: PrismaClient, *, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + clock: Callable[[], datetime] = get_utc_datetime, ): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client self._sleep = sleep + self._clock = clock @staticmethod def _get_prometheus_logger(): @@ -2260,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 @@ -2276,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) @@ -2287,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: @@ -2324,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: """ @@ -2387,12 +2503,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): where={ "file_purpose": "batch", "batch_processed": False, - "status": {"not_in": ["failed", "expired", "cancelled"]}, + "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]}, }, take=MAX_MATCHES_TO_RETURN, order={"created_at": "desc"}, ) + input_file_cutoff: Final = self._clock() - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) referencing_batches: Final[list[dict[str, object]]] = [] for batch in batches: try: @@ -2405,7 +2522,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): # Extract file IDs from batch # Batches typically reference the unified file ID in input_file_id # Output and error files are generated by the provider - input_file_id = batch_data.get("input_file_id") + input_file_id = batch_data.get("input_file_id") if batch.created_at >= input_file_cutoff else None output_file_id = batch_data.get("output_file_id") error_file_id = batch_data.get("error_file_id") 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/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/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/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/__init__.py b/litellm/__init__.py index b0ae899f838..a7f99829a9e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -682,6 +682,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 @@ -972,6 +973,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": @@ -1146,6 +1149,7 @@ model_list = list( | publicai_models | darkbloom_models | openinfer_models + | zerogpu_models | v0_models | morph_models | opencode_go_models @@ -1266,6 +1270,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, @@ -1454,6 +1459,7 @@ from .exceptions import ( Timeout, APIConnectionError, UnsupportedParamsError, + UnsupportedEndpointError as UnsupportedEndpointError, APIResponseValidationError, UnprocessableEntityError, InternalServerError, @@ -1743,6 +1749,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/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/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 98322aecf6a..47a3a903d82 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 @@ -43,8 +42,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, @@ -1745,19 +1744,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 2a20e812e68..1d5c6404bfd 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -129,6 +129,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 @@ -1096,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/exceptions.py b/litellm/exceptions.py index c8eee9ad1f8..cd5139f5f79 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -317,6 +317,10 @@ class MalformedRequestFieldError(BadRequestError): pass +class UnsupportedEndpointError(BadRequestError): + pass + + class UnprocessableEntityError(openai.UnprocessableEntityError): def __init__( self, @@ -1056,6 +1060,7 @@ LITELLM_EXCEPTION_TYPES: Final = [ BadRequestError, UnprocessableEntityError, UnsupportedParamsError, + UnsupportedEndpointError, Timeout, PermissionDeniedError, RateLimitError, diff --git a/litellm/images/main.py b/litellm/images/main.py index 3312ca417d5..84156919774 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -1062,7 +1062,7 @@ def image_edit( @client async def aimage_edit( - image: FileTypes | list[FileTypes], + image: FileTypes | list[FileTypes] | None, model: str, prompt: str, mask: str | None = None, @@ -1102,7 +1102,7 @@ async def aimage_edit( model=model, api_base=local_vars.get("base_url", None) ) - images: Final = image if isinstance(image, list) else [image] + images: Final = image if isinstance(image, list) else ([image] if image is not None else []) func: Final = partial( image_edit, diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index d8be607150c..e423b6698e6 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", 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/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/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/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 48e8a44b9a4..80f1f8f5eca 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) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 3c06c522dd6..cf6bdc69556 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 @@ -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/single_flight.py b/litellm/litellm_core_utils/single_flight.py new file mode 100644 index 00000000000..44a84d05925 --- /dev/null +++ b/litellm/litellm_core_utils/single_flight.py @@ -0,0 +1,92 @@ +import threading +from collections.abc import Callable +from concurrent.futures import Future +from dataclasses import dataclass, field +from functools import lru_cache +from typing import Final, Generic, ParamSpec, TypeVar + +_P: Final = ParamSpec("_P") +_R: Final = TypeVar("_R") + + +@dataclass(slots=True) +class SingleFlight(Generic[_R]): + lock: threading.Lock = field(default_factory=threading.Lock) + in_flight: Future[_R] | None = None + flights_started: int = 0 + last_flight_failed: bool = False + + def run(self, fetch: Callable[[], _R], while_busy: _R | None = None) -> _R: + leader: Final[Future[_R]] = Future() + with self.lock: + in_flight: Final = self.in_flight + joined_flight: Final = self.flights_started + retry_if_it_fails: Final = not self.last_flight_failed + if in_flight is None: + self._take_off(leader) + if in_flight is None: + return self._lead(leader, fetch) + if while_busy is not None: + return while_busy + try: + return in_flight.result() + except Exception: + retry: Final = self._claim_retry(joined_flight) if retry_if_it_fails else None + if retry is None: + raise + retry_flight, leads_retry = retry + return self._lead(retry_flight, fetch) if leads_retry else retry_flight.result() + + def _take_off(self, leader: Future[_R]) -> None: + self.in_flight = leader + self.flights_started += 1 + + def _claim_retry(self, failed_flight: int) -> tuple[Future[_R], bool] | None: + leader: Final[Future[_R]] = Future() + with self.lock: + in_flight: Final = self.in_flight + if in_flight is not None: + return in_flight, False + if self.flights_started > failed_flight and self.last_flight_failed: + return None + self._take_off(leader) + return leader, True + + def _lead(self, flight: Future[_R], fetch: Callable[[], _R]) -> _R: + try: + result: Final = fetch() + except BaseException as error: + self._land(failed=True) + flight.set_exception(error) + raise + self._land(failed=False) + flight.set_result(result) + return result + + def _land(self, failed: bool) -> None: + with self.lock: + self.in_flight = None + self.last_flight_failed = failed + + +@dataclass(frozen=True, slots=True) +class SingleFlightCache(Generic[_P, _R]): + cached: Callable[..., _R] + clear: Callable[[], None] + lock: threading.Lock = field(default_factory=threading.Lock) + + def __call__( + self, + *args: _P.args, + **kwargs: _P.kwargs, # kwargs-ok: forwarded verbatim to the cached builder, typed by its own signature + ) -> _R: + with self.lock: + return self.cached(*args, **kwargs) + + def cache_clear(self) -> None: + self.clear() + + +def single_flight_lru_cache(build: Callable[_P, _R]) -> SingleFlightCache[_P, _R]: + cached: Final = lru_cache(maxsize=128)(build) + return SingleFlightCache(cached, cached.cache_clear) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b64e146ca08..3fb25de43cc 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1557,6 +1557,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, 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/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index eb8e595c849..0967c22f93a 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -161,11 +161,58 @@ def _has_non_chat_user_block(message: Mapping[str, object]) -> bool: return any(isinstance(block, dict) and block.get("type") not in ValidUserMessageContentTypes for block in blocks) +_OPENAI_RESPONSES_BUILT_IN_TOOL_TYPES: Final = frozenset( + ( + "apply_patch", + "code_interpreter", + "computer", + "computer_use_preview", + "file_search", + "image_generation", + "local_shell", + "mcp", + "shell", + "web_search", + "web_search_2025_08_26", + "web_search_preview", + "web_search_preview_2025_03_11", + ) +) + + +_CHAT_INFERENCE_TRANSLATED_TOOL_TYPES: Final = frozenset(("namespace", "tool_search")) + + +def _is_responses_built_in_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return isinstance(tool_type, str) and tool_type in _OPENAI_RESPONSES_BUILT_IN_TOOL_TYPES + + +def _is_chat_inference_translated_tool(tool: Mapping[str, object]) -> bool: + tool_type: Final = tool.get("type") + return isinstance(tool_type, str) and tool_type in _CHAT_INFERENCE_TRANSLATED_TOOL_TYPES + + +def _needs_tool_translation(tool: Mapping[str, object]) -> bool: + return _is_chat_tool(tool) or _is_chat_inference_translated_tool(tool) or _is_responses_built_in_tool(tool) + + def _gets_placeholder_name(tool: Mapping[str, object]) -> bool: from litellm.llms.anthropic.messages.adapters.transformation import is_provider_native_tool_dict name: Final = tool.get("name") - return not (_is_chat_tool(tool) or is_provider_native_tool_dict(tool) or (isinstance(name, str) and name.strip())) + return not ( + _is_chat_tool(tool) + or _is_chat_inference_translated_tool(tool) + or is_provider_native_tool_dict(tool) + or (isinstance(name, str) and name.strip()) + ) + + +def _chat_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for + + return (tool,) if _is_chat_inference_translated_tool(tool) else chat_completion_tools_for(tool) def _untranslatable_input( @@ -182,8 +229,8 @@ def _untranslatable_input( unnamed_tool: Final = next(filter(_gets_placeholder_name, tools), None) if unnamed_tool is not None: return ChatCountTranslationError( - message=f"Invalid request for token counting: tool type {unnamed_tool.get('type')!r} " - "has no Chat Completions equivalent" + message=f"Invalid request for token counting: LiteLLM cannot count a tool of type " + f"{unnamed_tool.get('type')!r} for this model" ) return None @@ -223,9 +270,7 @@ def _translated_chat_request( chat_params: Mapping[str, object], translator: _ChatRequestTranslator, ) -> _ChatParams: - from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for - - chat_tools: Final = [dict(chat_tool) for chat_tool in chain.from_iterable(map(chat_completion_tools_for, tools))] + chat_tools: Final = [dict(chat_tool) for chat_tool in chain.from_iterable(map(_chat_tools_for, tools))] openai_params: Final = {name: value for name, value in chat_params.items() if name != "drop_params"} mapped_params: Final = {**openai_params, **({"tools": chat_tools} if chat_tools else {})} chat_messages: Final = cast( # cast-ok: chat-shaped request messages, read the way chat inference reads them @@ -255,7 +300,7 @@ def anthropic_fields_for_chat_tool_use( counted_tools: Final = tools or () if ( not chat_params - and not any(map(_is_chat_tool, counted_tools)) + and not any(map(_needs_tool_translation, counted_tools)) and not any(map(_needs_chat_translation, messages)) ): return None 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/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 23c5aa35f4b..4bd19889bcc 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -5,7 +5,6 @@ import os import re from collections.abc import Awaitable, Callable, Mapping from contextlib import closing -from functools import lru_cache from types import MappingProxyType from typing import Final, Literal, NamedTuple, cast @@ -18,15 +17,18 @@ import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.constants import DEFAULT_MAX_RETRIES +from litellm.litellm_core_utils.single_flight import SingleFlight, single_flight_lru_cache from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, + single_flight_bearer_token_provider, ) from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams +from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType from litellm.utils import add_path_to_api_base azure_ad_cache: Final = DualCache() @@ -94,7 +96,7 @@ def process_azure_headers(headers: httpx.Headers | dict) -> dict: return {**llm_response_headers, **openai_headers} -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_entra_id_token_provider( tenant_id: str, client_id: str, @@ -107,33 +109,38 @@ def _cached_entra_id_token_provider( to expiry, so reusing the provider keeps one AAD round trip per token lifetime instead of one per request. """ - from azure.identity import ClientSecretCredential, get_bearer_token_provider + from azure.identity import ClientSecretCredential - return get_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) + return single_flight_bearer_token_provider(ClientSecretCredential(tenant_id, client_id, client_secret), scope) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_username_password_token_provider( client_id: str, azure_username: str, azure_password: str, scope: str, ) -> Callable[[], str]: - from azure.identity import UsernamePasswordCredential, get_bearer_token_provider + from azure.identity import UsernamePasswordCredential credential: Final = UsernamePasswordCredential( client_id=client_id, username=azure_username, password=azure_password, ) - return get_bearer_token_provider(credential, scope) + return single_flight_bearer_token_provider(credential, scope) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: return get_azure_ad_token_provider(azure_scope=scope) +@single_flight_lru_cache +def _cached_default_azure_credential_provider(scope: str) -> Callable[[], str]: + return get_azure_ad_token_provider(azure_scope=scope, azure_credential=AzureCredentialType.DefaultAzureCredential) + + def cached_azure_ad_token_refresh_provider(scope: str) -> Callable[[], str]: return _cached_azure_ad_token_refresh_provider(scope) @@ -240,6 +247,58 @@ def get_azure_ad_token_from_username_password( return token_provider +@single_flight_lru_cache +def _oidc_token_exchange(cache_key: str) -> SingleFlight[str]: + return SingleFlight() + + +def _cached_oidc_token(cache_key: str) -> str | None: + cached: Final = azure_ad_cache.get_cache(cache_key) + return cached if isinstance(cached, str) else None + + +def _exchange_oidc_token(cache_key: str, token_url: str, client_id: str, scope: str, oidc_token: str) -> str: + cached_access_token: Final = _cached_oidc_token(cache_key) + if cached_access_token is not None: + return cached_access_token + + with closing(HTTPHandler(timeout=litellm.request_timeout)) as client: + req_token: Final = client.post( + token_url, + data={ + "client_id": client_id, + "grant_type": "client_credentials", + "scope": scope, + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": oidc_token, + }, + ) + + if req_token.status_code != 200: + raise AzureOpenAIError( + status_code=req_token.status_code, + message=req_token.text, + ) + + azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json() + azure_ad_token_access_token: Final = azure_ad_token_json.get("access_token", None) + azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None) + + if azure_ad_token_access_token is None: + raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned") + + if azure_ad_token_expires_in is None: + raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned") + + azure_ad_cache.set_cache( + key=cache_key, + value=azure_ad_token_access_token, + ttl=azure_ad_token_expires_in, + ) + + return azure_ad_token_access_token + + def get_azure_ad_token_from_oidc( azure_ad_token: str, azure_client_id: str | None = None, @@ -282,50 +341,25 @@ def get_azure_ad_token_from_oidc( "azure_client_id": azure_client_id, "azure_tenant_id": azure_tenant_id, "azure_authority_host": azure_authority_host, + "scope": scope, "oidc_token": oidc_token, } ) - azure_ad_token_access_token = azure_ad_cache.get_cache(azure_ad_token_cache_key) - if azure_ad_token_access_token is not None: - return azure_ad_token_access_token + cached_access_token: Final = _cached_oidc_token(azure_ad_token_cache_key) + if cached_access_token is not None: + return cached_access_token - with closing(HTTPHandler(timeout=litellm.request_timeout)) as client: - req_token: Final = client.post( - f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", - data={ - "client_id": azure_client_id, - "grant_type": "client_credentials", - "scope": scope, - "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", - "client_assertion": oidc_token, - }, + return _oidc_token_exchange(azure_ad_token_cache_key).run( + lambda: _exchange_oidc_token( + azure_ad_token_cache_key, + token_url=f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token", + client_id=azure_client_id, + scope=scope, + oidc_token=oidc_token, ) - - if req_token.status_code != 200: - raise AzureOpenAIError( - status_code=req_token.status_code, - message=req_token.text, - ) - - azure_ad_token_json: Final[_AzureAdTokenJson] = req_token.json() - azure_ad_token_access_token = azure_ad_token_json.get("access_token", None) - azure_ad_token_expires_in: Final = azure_ad_token_json.get("expires_in", None) - - if azure_ad_token_access_token is None: - raise AzureOpenAIError(status_code=422, message="Azure AD Token access_token not returned") - - if azure_ad_token_expires_in is None: - raise AzureOpenAIError(status_code=422, message="Azure AD Token expires_in not returned") - - azure_ad_cache.set_cache( - key=azure_ad_token_cache_key, - value=azure_ad_token_access_token, - ttl=azure_ad_token_expires_in, ) - return azure_ad_token_access_token - def select_azure_base_url_or_endpoint(azure_client_params: dict): azure_endpoint: Final = azure_client_params.get("azure_endpoint", None) @@ -503,17 +537,10 @@ class BaseAzureLLM(BaseOpenAILLM): Returns: Token provider callable if DefaultAzureCredential is enabled and available, None otherwise """ - from litellm.types.secret_managers.get_azure_ad_token_provider import ( - AzureCredentialType, - ) - verbose_logger.debug("Attempting to use DefaultAzureCredential for Azure Auth") try: - azure_ad_token_provider: Final = get_azure_ad_token_provider( - azure_scope=scope, - azure_credential=AzureCredentialType.DefaultAzureCredential, - ) + azure_ad_token_provider: Final = _cached_default_azure_credential_provider(scope) verbose_logger.debug("Successfully obtained Azure AD token provider using DefaultAzureCredential") return azure_ad_token_provider except Exception as e: 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/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/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 21fbdb04896..26bc5904e06 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -10,6 +10,7 @@ from typing import Any, Final from pydantic import JsonValue from typing_extensions import assert_never +import litellm from litellm._logging import verbose_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields @@ -53,6 +54,18 @@ class CountTokensFailure: CountTokensOutcome = CountedTokens | CountTokensFailure +def _claude_count_translator(litellm_params: Mapping[str, object]) -> AnthropicConfig: + deployment_model: Final = litellm_params.get("model") + if not isinstance(deployment_model, str): + return AnthropicConfig() + provider_hint: Final = litellm_params.get("custom_llm_provider") + model, provider, _, _ = litellm.get_llm_provider( + model=deployment_model, custom_llm_provider=provider_hint if isinstance(provider_hint, str) else None + ) + chat_config: Final = litellm.ProviderConfigManager.get_provider_chat_config(model, LlmProviders(provider)) + return chat_config if isinstance(chat_config, AnthropicConfig) else AnthropicConfig() + + def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) @@ -138,16 +151,22 @@ class BedrockTokenCounter(BaseTokenCounter): counted_model: Final = model_to_use.removeprefix("bedrock_mantle/") resolved_model: Final = get_bedrock_base_model(counted_model) primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE + litellm_params: Final = (deployment or {}).get("litellm_params", {}) translated: Final = ( await anthropic_count_fields( - counted_model, messages, system, tools, messages_format, AnthropicConfig(), chat_params + counted_model, + messages, + system, + tools, + messages_format, + _claude_count_translator(litellm_params), + chat_params, ) if is_mantle_claude_model(resolved_model) else None ) if isinstance(translated, ChatCountTranslationError): return translated.response(request_model, model_to_use, primary_tokenizer) - litellm_params: Final = (deployment or {}).get("litellm_params", {}) runtime_request: Final = ( _request_data(model_to_use, messages, tools, system) if translated is None 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/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/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/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_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/main.py b/litellm/main.py index 66c28372089..b8ad28c26fe 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6127,7 +6127,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) @@ -8394,7 +8405,7 @@ def transcription( custom_llm_provider=custom_llm_provider, ) - response: TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] | None = None + response: TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] provider_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( model=model, @@ -8556,10 +8567,16 @@ def transcription( provider_config=provider_config, shared_session=shared_session, ) + else: + raise litellm.UnsupportedEndpointError( + message=f"{custom_llm_provider}/{model} does not support audio transcription.", + model=model, + llm_provider=custom_llm_provider, + ) # Store duration in _hidden_params for cost calculation without # exposing it in the response body (see sync path comment above). - if response is not None and not isinstance(response, Coroutine): + if not isinstance(response, Coroutine): existing_duration: Final = getattr(response, "duration", None) if existing_duration is None: calculated_duration: Final = calculate_request_duration(file) @@ -8569,8 +8586,6 @@ def transcription( ) response_hidden_params["audio_transcription_duration"] = calculated_duration - if response is None: - raise ValueError("Unmapped provider passed in. Unable to get the response.") return response diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 02ee1e43285..95da1836aac 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, @@ -35598,6 +35800,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 +35884,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 +36067,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 +36125,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 +36247,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 +36303,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 +36351,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 +36445,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 +36538,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 +36639,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 +36731,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 +36821,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 +36913,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 +37004,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 +37250,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 +37305,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 +37375,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 +37437,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 +37496,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 +37563,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 +37628,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 +37685,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 +37745,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 +37807,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 +37869,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 +37928,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 +37988,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 +38042,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 +38091,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 +38148,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 +38265,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 +38354,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 +38412,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 +38467,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 +38522,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 +42319,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 +44295,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 +44410,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 +44496,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 +44540,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, @@ -65847,6 +66365,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 +73964,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": { @@ -76771,6 +77472,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 +77805,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 4451c8b0cd8..a614c117026 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, @@ -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..b10b5ae98b5 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -15015,7 +15015,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/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/common_request_processing.py b/litellm/proxy/common_request_processing.py index 24a36a7f2dc..f5d6417fb79 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 @@ -136,7 +137,7 @@ 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_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX @@ -2006,6 +2007,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 @@ -3498,6 +3542,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, @@ -4448,11 +4494,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: @@ -4523,10 +4573,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 @@ -4549,15 +4607,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/common_utils/token_count_settings.py b/litellm/proxy/common_utils/token_count_settings.py index 22b35bd1001..dd04e531f2d 100644 --- a/litellm/proxy/common_utils/token_count_settings.py +++ b/litellm/proxy/common_utils/token_count_settings.py @@ -13,6 +13,7 @@ from litellm.types.token_count import TokenCountProvider, TokenCountSourceFields from litellm.types.utils import TokenCountResponse CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404}) +REGION_DENIED_STATUSES: Final = frozenset({401, 403}) TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0 TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0 TOKENIZER_LOAD_MAX_RUNNING: Final = 4 @@ -47,6 +48,7 @@ class TokenCountModel: model: str region: str | None litellm_provider: TokenCountProvider | None = None + configured: bool = False @dataclass(frozen=True, slots=True) @@ -66,7 +68,13 @@ class ProviderCountFailure: def cannot_count_model(self) -> bool: return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES + @property + def denied_in_region(self) -> bool: + return not self.rejected_before_provider_call and self.status_code in REGION_DENIED_STATUSES + def with_hint(self, hint: str) -> "ProviderCountFailure": + if self.rejected_before_provider_call: + return self return replace(self, message=_with_hint(self.message, hint)) def to_proxy_exception(self) -> ProxyException: @@ -111,7 +119,7 @@ def _count_provider(value: object) -> TokenCountProvider | None: return None -def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: +def _count_model(value: Mapping[str, object], configured: bool) -> TokenCountModel | None: model: Final = _non_blank(value.get("model")) return ( None @@ -120,6 +128,7 @@ def _count_model(value: Mapping[str, object]) -> TokenCountModel | None: model=model, region=_non_blank(value.get("region")), litellm_provider=_count_provider(value.get("litellm_provider")), + configured=configured, ) ) @@ -146,12 +155,12 @@ def _resolved_count_model( ) -> TokenCountModel | None: configured_mapping: Final = as_str_mapping(configured) if configured_mapping is not None: - return _count_model(configured_mapping) + return _count_model(configured_mapping, configured=True) default_mapping: Final = as_str_mapping(default) return ( None if default_mapping is None - else _within_deployment_geography(_count_model(default_mapping), deployment_region) + else _within_deployment_geography(_count_model(default_mapping, configured=False), deployment_region) ) 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/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/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/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index a5a04fb68d2..9790c150ae1 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -242,6 +242,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. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7fcb77e2bbf..64fb2791913 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14853,6 +14853,7 @@ async def _try_native_responses_count( async def _provider_token_count( count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]], model_to_use: str, + request_model: str, settings: TokenCountSettings, local_estimate_allowed: bool, ) -> Union["TokenCountResponse", ProviderCountFailure, None]: @@ -14875,17 +14876,28 @@ async def _provider_token_count( if count_model is None: return counted.model_copy(update={"count_source": "provider"}) return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model}) - failure: Final = ( - counted - if count_model is None or count_model.region is None - else counted.with_hint( - f"This model counts tokens with {count_model.model} in {count_model.region}. " - "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" - ) + pinned_count_model: Final = None if count_model is None or count_model.region is None else count_model + region_hint: Final = ( + None + if pinned_count_model is None + else f"This model counts tokens with {pinned_count_model.model} in {pinned_count_model.region}. " + "If that region is not enabled for the AWS account, enable it or set token_count_model in model_info" ) - if count_model is None and failure.cannot_count_model and settings.framing is not None and local_estimate_allowed: - return failure - raise failure.to_proxy_exception() + failure: Final = counted if region_hint is None else counted.with_hint(region_hint) + framing_can_stand_in: Final = (count_model is None or not count_model.configured) and ( + failure.cannot_count_model or (pinned_count_model is not None and failure.denied_in_region) + ) + if not (framing_can_stand_in and settings.framing is not None and local_estimate_allowed): + raise failure.to_proxy_exception() + if pinned_count_model is not None: + verbose_proxy_logger.warning( + "Token count for %s answered %s from %s in %s, so it falls back to the framed local estimate", + request_model, + failure.status_code, + pinned_count_model.model, + pinned_count_model.region, + ) + return failure @dataclass(frozen=True, slots=True) @@ -15455,6 +15467,7 @@ async def _count_request_tokens( chat_params=MappingProxyType({}) if chat_count is None else chat_count.params, ), model_to_use=model_to_use, + request_model=request.model, settings=settings, local_estimate_allowed=litellm.disable_token_counter is not True and contents_estimable, ) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index a1820d831d3..423f129a9b4 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/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/responses/main.py b/litellm/responses/main.py index 8451834eab2..d474ed3991f 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -2341,6 +2341,7 @@ async def acancel_responses( timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, + litellm_metadata: Mapping[str, object] | None = None, **kwargs, ) -> ResponsesAPIResponse: """ @@ -2365,6 +2366,7 @@ async def acancel_responses( cancel_responses, response_id=response_id, custom_llm_provider=custom_llm_provider, + litellm_metadata=litellm_metadata, extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, @@ -2380,7 +2382,12 @@ async def acancel_responses( response = await init_response else: response = init_response - return response + + return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + responses_api_response=response, + litellm_metadata=litellm_metadata, + custom_llm_provider=custom_llm_provider, + ) except Exception as e: raise litellm.exception_type( model=None, @@ -2402,6 +2409,7 @@ def cancel_responses( timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, + litellm_metadata: Mapping[str, object] | None = None, **kwargs, ) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: """ @@ -2471,7 +2479,13 @@ def cancel_responses( shared_session=kwargs.get("shared_session"), ) - return response + if not isinstance(response, ResponsesAPIResponse): + return response + return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( + responses_api_response=response, + litellm_metadata=litellm_metadata, + custom_llm_provider=custom_llm_provider, + ) except Exception as e: raise litellm.exception_type( model=None, 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 a6b042a2234..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, @@ -399,7 +403,7 @@ class ResponsesAPIRequestUtils: def update_responses_api_response_id_with_model_id( responses_api_response: ResponsesAPIResponse, custom_llm_provider: str | None, - litellm_metadata: dict[str, object] | None = None, + litellm_metadata: Mapping[str, object] | None = None, ) -> ResponsesAPIResponse: ... @@ -408,7 +412,7 @@ class ResponsesAPIRequestUtils: def update_responses_api_response_id_with_model_id( responses_api_response: dict[str, object], custom_llm_provider: str | None, - litellm_metadata: dict[str, object] | None = None, + litellm_metadata: Mapping[str, object] | None = None, ) -> dict[str, object]: ... @@ -418,7 +422,7 @@ class ResponsesAPIRequestUtils: def update_responses_api_response_id_with_model_id( responses_api_response: ResponsesAPIResponse | dict[str, Any], custom_llm_provider: str | None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, Any] | None = None, ) -> ResponsesAPIResponse | dict[str, object]: """Update the responses_api_response_id with model_id and custom_llm_provider. @@ -1110,7 +1114,7 @@ class ResponsesAPIRequestUtils: def _update_container_ids_in_response( responses_api_response: ResponsesAPIResponse | dict[str, object], custom_llm_provider: str | None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, Any] | None = None, ) -> ResponsesAPIResponse | dict[str, object]: """Encode container IDs in the response output with provider/model info. @@ -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 f0bb6c3a7a9..5b18d36316f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -216,6 +216,8 @@ from litellm.router_utils.common_utils import ( deployments_the_caller_may_use, filter_team_based_models, filter_web_search_deployments, + format_content_policy_fallback_hint, + format_content_policy_no_fallback_note, format_fallback_outcome_message, format_no_fallback_group_message, get_request_team_id, @@ -631,6 +633,7 @@ _DEPLOYMENT_PARAMS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | GenericLite _OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) _DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" _CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) +_DEPLOYMENTS_ADAPTER: Final = TypeAdapter(Sequence[Mapping[str, object]]) def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: @@ -692,6 +695,12 @@ def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, obj return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) +def _every_deployment_refused( + deployments: Sequence[Mapping[str, object]], refused_deployment_ids: tuple[str, ...] +) -> bool: + return all(_configured_model_info(deployment).get("id") in refused_deployment_ids for deployment in deployments) + + def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]: if silent_model is None: return () @@ -832,6 +841,10 @@ def _request_model_group(kwargs: Mapping[str, object]) -> str | None: return model_group if isinstance(model_group, str) else None +def _mid_stream_error_before_output(error: "MidStreamFallbackError") -> bool: + return not error.generated_content or error.is_pre_first_chunk + + def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception: """The provider's own error, which is what the retry policy and should_retry_this_error classify.""" return error.original_exception if error.original_exception is not None else error @@ -962,6 +975,10 @@ def _responses_stream_holds_event(item: object, held_event_count: int) -> bool: return getattr(item, "type", None) in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES +def _responses_event_announces_background_run(item: object) -> bool: + return getattr(getattr(item, "response", None), "background", None) is True + + class FallbackAwareAnthropicMessagesStream: """ Bare async generators can't carry the `_hidden_params` attribute the @@ -4127,11 +4144,16 @@ class Router: async def stream_with_fallbacks(): held_lifecycle_events: tuple[object, ...] = () # rebind-ok: flushed at first output, dropped on fallback + response_id_exposed = False # rebind-ok: set once, before an early-released lifecycle event is yielded try: async for item in cast( # cast-ok: response iterators are async streams AsyncIterator[object], source_iterator ): - if _responses_stream_holds_event(item, len(held_lifecycle_events)): + if not response_id_exposed and self._aresponses_background_stream_can_release( + item, source_iterator, initial_kwargs + ): + response_id_exposed = True + if not response_id_exposed and _responses_stream_holds_event(item, len(held_lifecycle_events)): held_lifecycle_events = (*held_lifecycle_events, item) continue for held_event in held_lifecycle_events: @@ -4141,6 +4163,8 @@ class Router: for held_event in held_lifecycle_events: yield held_event except MidStreamFallbackError as e: + if response_id_exposed: + self._raise_released_stream_error(e, initial_kwargs) async with contextlib.aclosing( self._aresponses_fallback_attempt( e, source_iterator, initial_kwargs, wrapper.adopt_fallback_headers, held_lifecycle_events @@ -4166,6 +4190,54 @@ class Router: wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) return wrapper + def _aresponses_background_stream_can_release( + self, + item: object, + source_iterator: "BaseResponsesAPIStreamingIterator", + kwargs: Mapping[str, object], + ) -> bool: + from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator + + if ( + type(source_iterator) is not ResponsesAPIStreamingIterator + or source_iterator.custom_llm_provider != "openai" + ): + return False + if not _responses_event_announces_background_run(item): + return False + model_group: Final = kwargs.get("model") + return isinstance(model_group, str) and not self._stream_can_fall_back(model_group, kwargs) + + def _raise_released_stream_error( + self, e: "MidStreamFallbackError", initial_kwargs: Mapping[str, object] + ) -> NoReturn: + original: Final = e.original_exception + model_group: Final = initial_kwargs.get("model") + if ( + litellm.expose_router_debug_in_errors + and isinstance(original, litellm.ContentPolicyViolationError) + and isinstance(model_group, str) + and not fallbacks_disabled_for_request(initial_kwargs) + and _mid_stream_error_before_output(e) + ): + hop_depth: Final = initial_kwargs.get("fallback_depth") + original.message += format_content_policy_no_fallback_note( + model_group, + fallback_lookup_groups(initial_kwargs, model_group), + _request_fallback_list( + initial_kwargs, + "content_policy_fallbacks", + cast("list[object] | None", self.content_policy_fallbacks), # cast-ok: untyped router attribute + ), + _request_fallback_list( + initial_kwargs, + "fallbacks", + cast("list[object] | None", self.fallbacks), # cast-ok: untyped router attribute + ), + isinstance(hop_depth, int) and hop_depth > 0, + ) + Router._raise_original_mid_stream_error(e) + def _responses_continuation( self, e: "MidStreamFallbackError", @@ -4175,7 +4247,7 @@ class Router: """What a fallback continues from once a Responses stream broke after text, raising the primary's error when the request constrains its output or the stream began a tool call, since a continued JSON document or tool call would not parse, and when a content filter cut it, so no other model finishes filtered text.""" - if not e.generated_content or e.is_pre_first_chunk: + if _mid_stream_error_before_output(e): return None if source_iterator.started_tool_call or isinstance(e.original_exception, litellm.ContentPolicyViolationError): Router._raise_original_mid_stream_error(e) @@ -4706,7 +4778,13 @@ class Router: @staticmethod def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None: effective_model_info: Final = kwargs.get("model_info") - deployment_id: Final = effective_model_info.get("id") if isinstance(effective_model_info, Mapping) else None + if not isinstance(effective_model_info, Mapping): + return + model_info: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info) + configured_deployment_id: Final = ( + model_info.get("original_model_id") if isinstance(exception, litellm.UnsupportedEndpointError) else None + ) + deployment_id: Final = configured_deployment_id or model_info.get("id") if isinstance(deployment_id, str) and deployment_id: exception.retry_skip_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # dynamic stamp, read by _deployment_ids_to_skip_on_retry @@ -6378,6 +6456,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 ) @@ -6487,7 +6568,7 @@ class Router: model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group has_generated_content = not ( # rebind-ok: set once real content is seen, the buffer cap is hit, or neither a retry nor a fallback can take over self._anthropic_messages_stream_can_retry(initial_kwargs) - or self._anthropic_messages_stream_can_fall_back(model, initial_kwargs) + or self._stream_can_fall_back(model, initial_kwargs) ) buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline try: @@ -6710,7 +6791,7 @@ class Router: budget, _ = self._stream_retry_budget(frame_error, kwargs) if budget > attempted_retries_for_request(kwargs): return frame_error - if self._anthropic_messages_stream_can_fall_back(model_group, kwargs): + if self._stream_can_fall_back(model_group, kwargs): return frame_error return None @@ -8725,7 +8806,9 @@ class Router: ) return response else: - error_message = f"model={model_group}. content_policy_fallback={mask_sensitive_structure(content_policy_fallbacks)}. fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'content_policy_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" + error_message = format_content_policy_fallback_hint( + model_group, content_policy_fallbacks, fallbacks + ) verbose_router_logger.info( msg=f"Got 'ContentPolicyViolationError'. No content_policy_fallback set. Defaulting \ to fallbacks, if available.{error_message}" @@ -9048,6 +9131,10 @@ class Router: ) if first_skipped_ids: kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it + if isinstance(original_exception, litellm.UnsupportedEndpointError) and _every_deployment_refused( + _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), first_skipped_ids + ): + raise else: raise @@ -9127,6 +9214,10 @@ class Router: ) if skipped_ids: kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it + if isinstance(e, litellm.UnsupportedEndpointError) and _every_deployment_refused( + _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), skipped_ids + ): + raise _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, @@ -9223,7 +9314,7 @@ class Router: status_code: Final = getattr(error, "status_code", None) if status_code is not None and not litellm.should_retry(status_code): # 401/403 are special cases - allow retry if multiple deployments exist (handled below) - if status_code not in (401, 403): + if status_code not in (401, 403) and not isinstance(error, litellm.UnsupportedEndpointError): raise error if isinstance(error, litellm.NotFoundError): @@ -9873,7 +9964,7 @@ class Router: ) return has_unattempted_fallback_target(resolved, kwargs) - def _anthropic_messages_order_levels(self, model_group: str, kwargs: Mapping[str, Any]) -> tuple[int, ...]: + def _stream_order_levels(self, model_group: str, kwargs: Mapping[str, Any]) -> tuple[int, ...]: """ The distinct deployment order levels the fallback dispatcher would see for this request, computed the same way: the tier a pre-routing hook selected wins over the requested group. @@ -9891,7 +9982,7 @@ class Router: ) ) - def _anthropic_messages_stream_can_fall_back(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: + def _stream_can_fall_back(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: """ Whether async_function_with_fallbacks_common_utils could still route a MidStreamFallbackError somewhere for this request (order levels, weighted @@ -9903,7 +9994,7 @@ class Router: return False if self.enable_weighted_failover: return True - order_levels: Final = self._anthropic_messages_order_levels(model_group, kwargs) + order_levels: Final = self._stream_order_levels(model_group, kwargs) if len(order_levels) > 1: current_target: Final = kwargs.get("_target_order") skip_up_to: Final = current_target if current_target is not None else order_levels[0] @@ -10276,28 +10367,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( @@ -13857,7 +13958,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, @@ -13870,6 +13973,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 () @@ -14439,7 +14548,6 @@ class Router: Filter out model in model group, if: - model context window < message length. For azure openai models, requires 'base_model' is set. - https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models - - filter models above rpm limits - if region given, filter out models not in that region / unknown region - [TODO] function call and model doesn't support function calling """ @@ -14454,8 +14562,6 @@ class Router: _context_window_error = False _potential_error_str = "" - _rate_limit_error = False - parent_otel_span: Final = get_parent_otel_span_from_kwargs(request_kwargs) has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending( request_kwargs @@ -14479,13 +14585,6 @@ class Router: else frozenset() ) - ## get model group RPM ## - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") - rpm_key: Final = f"{model}:rpm:{current_minute}" - model_group_cache: Final = ( - self.cache.get_cache(key=rpm_key, local_only=True, parent_otel_span=parent_otel_span) or {} - ) # check the in-memory cache used by lowest_latency and usage-based routing. Only check the local cache. for idx, deployment in enumerate(_returned_deployments): # Cache nested dict access to avoid repeated temporary dict allocations _litellm_params = deployment.get("litellm_params", {}) @@ -14508,24 +14607,6 @@ class Router: ) continue - model_id = _model_info.get("id", "") - ## RPM CHECK ## - ### get local router cache ### - current_request_cache_local = ( - self.cache.get_cache(key=model_id, local_only=True, parent_otel_span=parent_otel_span) or 0 - ) - ### get usage based cache ### - if isinstance(model_group_cache, dict) and self.routing_strategy != "usage-based-routing-v2": - model_group_cache[model_id] = model_group_cache.get(model_id, 0) - - current_request = max(current_request_cache_local, model_group_cache[model_id]) - - if isinstance(_litellm_params, dict) and _litellm_params.get("rpm", None) is not None: - if isinstance(_litellm_params["rpm"], int) and _litellm_params["rpm"] <= current_request: - invalid_model_indices.add(idx) - _rate_limit_error = True - continue - ## REGION CHECK ## if request_kwargs is not None and request_kwargs.get("allowed_model_region") is not None: allowed_model_region = request_kwargs.get("allowed_model_region") @@ -14562,12 +14643,6 @@ class Router: ) if len(excluded_indices) == len(_returned_deployments): - """ - - no healthy deployments available b/c context window checks or rate limit error - - - First check for rate limit errors (if this is true, it means the model passed the context window check but failed the rate limit check) - """ - group_deployments: Final = ( tuple(request_deployments) if request_deployments is not None @@ -14589,12 +14664,7 @@ class Router: ), model=model, ) - if _rate_limit_error is True: # allow generic fallback logic to take place - raise RouterRateLimitErrorBasic( - model=model, - ) - - elif _context_window_error is True: + if _context_window_error is True: raise litellm.ContextWindowExceededError( message=f"litellm._pre_call_checks: Context Window exceeded for given call. No models have context window large enough for this call.\n{_potential_error_str}", model=model, diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 91b9e5cb114..2451422b5e6 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -39,6 +39,9 @@ _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = tuple[Mapping[str, object], ...] ) _DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) +_FALLBACK_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( + tuple[Mapping[str, object], ...] +) def with_admin_api_key_secret_reference(deployment: Deployment) -> Deployment: @@ -256,6 +259,40 @@ def format_fallback_outcome_message( return f"{lead} Fallback model group(s) configured: {targets}." +def format_content_policy_fallback_hint( + model_group: str | None, content_policy_fallbacks: object, fallbacks: object +) -> str: + return ( + f"model={model_group}. content_policy_fallback={mask_sensitive_structure(content_policy_fallbacks)}. " + f"fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'content_policy_fallback' - " + "https://docs.litellm.ai/docs/routing#fallbacks" + ) + + +def _as_fallback_entries(fallbacks: object) -> tuple[Mapping[str, object], ...] | None: + try: + return _FALLBACK_ENTRIES_ADAPTER.validate_python(fallbacks) + except ValidationError: + return None + + +def format_content_policy_no_fallback_note( + model_group: str, + lookup_groups: tuple[str, ...], + content_policy_fallbacks: object, + fallbacks: object, + nested_fallback_hop: bool, +) -> str: + outcome: Final = "" if nested_fallback_hop else format_fallback_outcome_message(model_group, None, "") + if content_policy_fallbacks is not None: + return outcome + hint: Final = f"\n{format_content_policy_fallback_hint(model_group, content_policy_fallbacks, fallbacks)}" + fallback_entries: Final = _as_fallback_entries(fallbacks) if fallbacks is not None else None + if fallback_entries is None or not lookup_groups or nested_fallback_hop: + return hint + outcome + return hint + format_no_fallback_group_message(lookup_groups, fallback_entries) + + def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str: """ Hash of the credential params, used for mapping the file id to the right model diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index e1baf66b201..128bca6cca0 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -356,7 +356,7 @@ def _should_run_cooldown_logic( ) return False - if isinstance(original_exception, litellm.MalformedRequestFieldError): + if isinstance(original_exception, (litellm.MalformedRequestFieldError, litellm.UnsupportedEndpointError)): verbose_router_logger.debug( "Should Not Run Cooldown Logic: litellm rejected the request before calling the deployment" ) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 6c93fe6d239..0a77ad5b9ea 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -261,7 +261,10 @@ def is_bad_request_without_fallback(exception: Exception, fallback_on_unsupporte return False if isinstance(exception, litellm.UnsupportedParamsError): return True - if isinstance(exception, (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError)): + if isinstance( + exception, + (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError, litellm.UnsupportedEndpointError), + ): return False return exception.status_code == 400 and bool(getattr(exception, "llm_provider", "")) diff --git a/litellm/router_utils/get_retry_from_policy.py b/litellm/router_utils/get_retry_from_policy.py index 7c412dc759b..835cfaacfd5 100644 --- a/litellm/router_utils/get_retry_from_policy.py +++ b/litellm/router_utils/get_retry_from_policy.py @@ -13,6 +13,7 @@ from litellm.exceptions import ( RateLimitError, ServiceUnavailableError, Timeout, + UnsupportedEndpointError, ) from litellm.types.router import RetryPolicy @@ -59,6 +60,8 @@ def get_num_retries_from_retry_policy( policy: Final = resolve_retry_policy(retry_policy, model_group, model_group_retry_policy) if policy is None: return None + if isinstance(exception, UnsupportedEndpointError): + return policy.DefaultRetries by_class: Final = ( _RETRIES_BY_EXCEPTION_TYPE[cls](policy) for cls in type(exception).__mro__ if cls in _RETRIES_BY_EXCEPTION_TYPE ) 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/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index 2b707f179d0..4967e199218 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -1,14 +1,27 @@ import os +import time from collections.abc import Callable +from dataclasses import dataclass +from types import TracebackType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.single_flight import SingleFlight from litellm.types.secret_managers.get_azure_ad_token_provider import ( AzureCredentialType, ) +_HELD_TOKEN_EXPIRY_MARGIN_SECONDS: Final = 60 + if TYPE_CHECKING: - from azure.core.credentials import TokenProvider + from azure.core.credentials import ( + AccessToken, + AccessTokenInfo, + SupportsTokenInfo, + TokenCredential, + TokenRequestOptions, + ) + from typing_extensions import TypeIs # noqa: TID251 # tells the two azure-core credential protocols apart def infer_credential_type_from_environment() -> AzureCredentialType: @@ -37,6 +50,86 @@ def infer_credential_type_from_environment() -> AzureCredentialType: return AzureCredentialType.DefaultAzureCredential +@dataclass(slots=True) +class _HeldToken: + token: "AccessToken | AccessTokenInfo | None" = None + + def still_valid(self) -> str | None: + held: Final = self.token + if held is None or held.expires_on - time.time() <= _HELD_TOKEN_EXPIRY_MARGIN_SECONDS: + return None + return held.token + + +@dataclass(slots=True) +class _TokenInfoWatcher: + credential: "SupportsTokenInfo" + held: _HeldToken + + def get_token_info(self, *scopes: str, options: "TokenRequestOptions | None" = None) -> "AccessTokenInfo": + token: Final = self.credential.get_token_info(*scopes, options=options) + self.held.token = token + return token + + def close(self) -> None: + self.credential.close() + + def __enter__(self) -> "SupportsTokenInfo": + return self.credential.__enter__() + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.credential.__exit__(exc_type, exc_value, traceback) + + +@dataclass(slots=True) +class _TokenWatcher: + credential: "TokenCredential" + held: _HeldToken + + def get_token( + self, + *scopes: str, + claims: str | None = None, + tenant_id: str | None = None, + enable_cae: bool = False, + **kwargs: object, # kwargs-ok: azure-core forwards its request options verbatim to the wrapped credential + ) -> "AccessToken": + token: Final = self.credential.get_token( + *scopes, claims=claims, tenant_id=tenant_id, enable_cae=enable_cae, **kwargs + ) + self.held.token = token + return token + + +def _supports_token_info( + credential: "TokenCredential | SupportsTokenInfo", +) -> "TypeIs[SupportsTokenInfo]": + return hasattr(credential, "get_token_info") + + +def single_flight_bearer_token_provider( + credential: "TokenCredential | SupportsTokenInfo", scope: str +) -> Callable[[], str]: + from azure.identity import get_bearer_token_provider + + held: Final = _HeldToken() + watched: Final = ( + _TokenInfoWatcher(credential, held) if _supports_token_info(credential) else _TokenWatcher(credential, held) + ) + provider: Final = get_bearer_token_provider(watched, scope) + token_request: Final[SingleFlight[str]] = SingleFlight() + + def bearer_token() -> str: + return token_request.run(provider, while_busy=held.still_valid()) + + return bearer_token + + def get_azure_ad_token_provider( azure_scope: str | None = None, azure_credential: AzureCredentialType | None = None, @@ -65,7 +158,6 @@ def get_azure_ad_token_provider( DefaultAzureCredential, ManagedIdentityCredential, WorkloadIdentityCredential, - get_bearer_token_provider, ) if azure_scope is None: @@ -77,7 +169,7 @@ def get_azure_ad_token_provider( else None or os.environ.get("AZURE_CREDENTIAL") or infer_credential_type_from_environment() ) verbose_logger.info("For Azure AD Token Provider, choosing credential type: %s", cred) - credential: TokenProvider | None = None + credential: TokenCredential | SupportsTokenInfo | None = None if cred == AzureCredentialType.ClientSecretCredential: credential = ClientSecretCredential( client_id=os.environ["AZURE_CLIENT_ID"], @@ -133,4 +225,4 @@ def get_azure_ad_token_provider( if credential is None: raise ValueError("No credential provided") - return get_bearer_token_provider(credential, azure_scope) + return single_flight_bearer_token_provider(credential, azure_scope) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index b33cf0e6b7e..588bdb1409c 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -5,7 +5,6 @@ import time import traceback from collections.abc import Callable from dataclasses import dataclass -from functools import lru_cache from typing import Final import httpx @@ -14,6 +13,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.single_flight import single_flight_lru_cache from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, @@ -27,7 +27,7 @@ _PARSED_LITERAL: Final = TypeAdapter(object) _SECRET_MANAGER_CLIENT: Final = TypeAdapter(object) -@lru_cache(maxsize=128) +@single_flight_lru_cache def _cached_azure_oidc_token_provider(audience: str) -> Callable[[], str]: return get_azure_ad_token_provider(azure_scope=audience) 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/utils.py b/litellm/types/utils.py index 38b4c3de327..e4d68a68223 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 @@ -4587,6 +4589,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 a232fd2f96a..29552c2667b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7001,6 +7001,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), @@ -9620,18 +9622,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 diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 02ee1e43285..95da1836aac 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, @@ -35598,6 +35800,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 +35884,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 +36067,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 +36125,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 +36247,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 +36303,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 +36351,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 +36445,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 +36538,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 +36639,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 +36731,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 +36821,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 +36913,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 +37004,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 +37250,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 +37305,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 +37375,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 +37437,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 +37496,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 +37563,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 +37628,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 +37685,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 +37745,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 +37807,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 +37869,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 +37928,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 +37988,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 +38042,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 +38091,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 +38148,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 +38265,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 +38354,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 +38412,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 +38467,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 +38522,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 +42319,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 +44295,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 +44410,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 +44496,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 +44540,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, @@ -65847,6 +66365,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 +73964,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": { @@ -76771,6 +77472,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 +77805,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 4451c8b0cd8..a614c117026 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, @@ -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/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/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index f4a295e697d..eee02a2f6b6 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -95,6 +95,7 @@ ignored_function_names = [ "arm_routing_read_prefetch", # Tested in tests/unit/caching/test_request_redis_batch_pre_call.py (file lacks "router" in name) "_deployment_params_base_model", # Tested through log_retry in test_router_helper_utils.py "_configured_model_info", # Tested through get_configured_service_tiers in test_router.py + "_every_deployment_refused", # Tested through the unsupported transcription retry tests in test_router.py "_routable_deployments", # Tested through get_configured_service_tiers and get_routable_upstream_model in test_router.py "_async_get_available_deployment", # Body of the `route {model}` phase wrapper, exercised through async_get_available_deployment in test_router.py "_async_get_available_deployment_for_pass_through", # Same, through async_get_available_deployment_for_pass_through in test_router.py @@ -174,6 +175,10 @@ 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 + "_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 + "_mid_stream_error_before_output", # Tested through the Responses mid-stream fallback and background stream tests ] 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/test_accounting_ledger.py b/tests/integration/database/test_accounting_ledger.py new file mode 100644 index 00000000000..089ed872dee --- /dev/null +++ b/tests/integration/database/test_accounting_ledger.py @@ -0,0 +1,128 @@ +import os +import subprocess +import sys +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import psycopg +import pytest +from psycopg import sql + +from litellm.accounting import DeliveryState, DeliveryTarget, EventKind +from tests.integration._support.database import scratch_database + +_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) 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_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index 27cbbe81a13..380881a6796 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -37,6 +37,7 @@ from integration.providers._count_tokens_system_lift import ( REMINDER, RESPONSES_REQUEST_OPTIONS, RESPONSES_TOOL_USE, + RESPONSES_TOOLS, STRING_CASE, TOKEN_COUNTING_BETA, TOOLS, @@ -603,7 +604,24 @@ def _generating(provider: _Provider) -> Callable[[Request], Reply]: return respond -@pytest.mark.parametrize("request_body", RESPONSES_TOOL_USE.values(), ids=RESPONSES_TOOL_USE.keys()) +_REWRITTEN_TOOL_USE: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + **RESPONSES_TOOL_USE, + "file-search": { + "input": WEATHER_QUESTION, + "tools": [{"type": "file_search", "vector_store_ids": ["vs_audit_weather"]}], + }, + "additional-tools": { + "input": [ + {"role": "user", "content": WEATHER_QUESTION}, + {"type": "additional_tools", "tools": RESPONSES_TOOLS}, + ], + }, + } +) + + +@pytest.mark.parametrize("request_body", _REWRITTEN_TOOL_USE.values(), ids=_REWRITTEN_TOOL_USE.keys()) def test_responses_input_tokens_counts_what_responses_inference_sends( deployment: _Deployment, gateway: Gateway, request_body: dict[str, JsonValue] ) -> None: @@ -642,6 +660,79 @@ def test_responses_input_tokens_counts_the_request_options_responses_inference_s assert count_request_fields(counted) == inference_count_fields(inference), (counted, inference) +def _without_cache_marks(value: JsonValue) -> JsonValue: + if isinstance(value, dict): + return {key: _without_cache_marks(item) for key, item in value.items() if key != "cache_control"} + if isinstance(value, list): + return [_without_cache_marks(item) for item in value] + return value + + +def test_responses_input_tokens_leaves_out_the_cache_marks_responses_inference_sends( + deployment: _Deployment, gateway: Gateway +) -> None: + marked: Final[dict[str, JsonValue]] = {"type": "ephemeral"} + request_body: Final[dict[str, JsonValue]] = { + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": f"Weather note {index}.", "cache_control": marked} + for index in range(2) + ], + }, + {"role": "user", "content": WEATHER_QUESTION}, + ], + "tools": RESPONSES_TOOLS, + } + with wire_server(_generating(deployment.provider), port=deployment.port) as peer: + generated: Final = gateway.request( + "POST", "/v1/responses", {**request_body, "model": deployment.model, "max_output_tokens": 16} + ) + assert generated.status_code == 200, generated.text + response: Final = gateway.request( + "POST", "/v1/responses/input_tokens", {**request_body, "model": deployment.model} + ) + assert _payload(response)["input_tokens"] == _COUNT, response.text + inference, counted = (counted_fields(body) for body in _inference_and_count(peer)) + assert "cache_control" in json.dumps(inference), inference + assert "cache_control" not in json.dumps(counted), counted + assert counted == _without_cache_marks(inference), (counted, inference) + + +_UNCOUNTABLE_TOOLS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + { + "computer_use_preview": { + "type": "computer_use_preview", + "display_width": 1024, + "display_height": 768, + "environment": "browser", + }, + "mcp": { + "type": "mcp", + "server_label": "docs", + "server_url": "http://127.0.0.1:9/sse", + "require_approval": "never", + }, + } +) + + +@pytest.mark.parametrize("tool", _UNCOUNTABLE_TOOLS.values(), ids=_UNCOUNTABLE_TOOLS.keys()) +def test_responses_input_tokens_rejects_a_tool_it_cannot_count_without_calling_the_peer( + deployment: _Deployment, gateway: Gateway, tool: dict[str, JsonValue] +) -> None: + with _listening(deployment) as peer: + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": deployment.model, "input": WEATHER_QUESTION, "tools": [tool]}, + ) + assert response.status_code == 400, response.text + assert tool["type"] in response.text, response.text + assert peer.drain() == () + + def _text_blocks(message: JsonValue) -> JsonValue: if not isinstance(message, dict) or not isinstance(message.get("content"), str): return message diff --git a/tests/integration/providers/test_azure_entra_token_reuse_wire.py b/tests/integration/providers/test_azure_entra_token_reuse_wire.py new file mode 100644 index 00000000000..85aef110124 --- /dev/null +++ b/tests/integration/providers/test_azure_entra_token_reuse_wire.py @@ -0,0 +1,625 @@ +import asyncio +import itertools +import json +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from enum import StrEnum +from pathlib import Path +from queue import SimpleQueue +from typing import Final +from urllib.parse import parse_qs, urlsplit + +import httpx +import pytest +import websockets +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.ports import allocate_port +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_ws_peer import Peer, completed, responses_peer, session +from integration._support.tls import write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue + +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] + +_DEPLOYMENT: Final = "gpt-5.4-nano" +_DATED_API_VERSION: Final = "2024-10-21" +_REALTIME_API_VERSION: Final = "2024-10-01-preview" +_IDENTITY_PATH: Final = "/msi/token" +_EXCHANGE_SUFFIX: Final = "/oauth2/v2.0/token" +_FOUNDRY_PROJECT: Final = "/api/projects/audit" +_BING_TOOL: Final = "audit-bing-entra" +_BING_RESOURCE: Final = "https://ai.azure.com" +_BING_CITATION: Final = "https://result.invalid/bing" +_TENANT: Final = "audit-tenant" +_APP_CLIENT: Final = "audit-app-client" +_HOLD_SECONDS: Final = 60 +_LIVELINESS_SECONDS: Final = 3 +_HERD: Final = 6 +_HERD_WINDOW_SECONDS: Final = 3 +_TOKEN_LIFETIME_SECONDS: Final = 3600 +_INSIDE_REFRESH_WINDOW_SECONDS: Final = 120 +_REMOVED_ENVIRONMENT: Final = ( + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_TENANT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_FEDERATED_TOKEN_FILE", + "AZURE_USERNAME", + "AZURE_PASSWORD", + "AZURE_CREDENTIAL", + "AZURE_SCOPE", + "MSI_ENDPOINT", + "MSI_SECRET", + "IMDS_ENDPOINT", + "IDENTITY_SERVER_THUMBPRINT", + "BING_GROUNDING_TOKEN", + "HTTP_PROXY", + "HTTPS_PROXY", + "http_proxy", + "https_proxy", +) + + +class _Surface(StrEnum): + LEGACY_CHAT = "legacy-chat" + V1_CHAT = "v1-chat" + RESPONSES = "responses" + RESPONSES_WEBSOCKET = "responses-websocket" + REALTIME = "realtime" + REALTIME_HEALTH_CHECK = "realtime-health-check" + + +_HTTP_SURFACES: Final = frozenset({_Surface.LEGACY_CHAT, _Surface.V1_CHAT, _Surface.RESPONSES}) + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + port: int + cert: tuple[Path, Path] + + @property + def edge(self) -> str: + return f"http://127.0.0.1:{self.port}" + + +@dataclass(frozen=True, slots=True) +class _Hold: + resource: str + arrivals: SimpleQueue[str] = field(default_factory=SimpleQueue) + release: threading.Event = field(default_factory=threading.Event) + + +def _config(directory: Path) -> Path: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = directory / "azure-entra-token-reuse.yaml" + path.write_text( + yaml.safe_dump( + { + **stock, + "litellm_settings": {**object_value(stock["litellm_settings"]), "enable_azure_ad_token_refresh": True}, + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0}, + "search_tools": [ + {"search_tool_name": _BING_TOOL, "litellm_params": {"search_provider": "bing_grounding"}} + ], + } + ) + ) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("azure-entra-token-reuse") + port: Final = allocate_port() + cert: Final = write_self_signed_cert(directory) + edge: Final = f"http://127.0.0.1:{port}" + environment: Final = { + "AZURE_CLIENT_ID": "audit-managed-identity", + "IDENTITY_ENDPOINT": f"{edge}{_IDENTITY_PATH}", + "IDENTITY_HEADER": "audit-identity-header", + "AZURE_AUTHORITY_HOST": edge, + "BING_GROUNDING_PROJECT_ENDPOINT": f"{edge}{_FOUNDRY_PROJECT}", + "BING_GROUNDING_MODEL": _DEPLOYMENT, + "BING_GROUNDING_CONNECTION_ID": "audit-bing-connection", + "SSL_CERT_FILE": str(cert[0]), + } + with ( + gateway_from_environment() as shared, + owned_proxy( + shared, directory, environment, config=_config(directory), remove_environment=_REMOVED_ENVIRONMENT + ) as proxy, + ): + yield _Rig(proxy, port, cert) + + +def _json(payload: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps(payload).encode()) + + +def _identity_token(resource: str) -> str: + return f"mi-token:{resource}" + + +def _identity_resource(request: Request) -> str | None: + target: Final = urlsplit(request.target) + if target.path != _IDENTITY_PATH: + return None + return parse_qs(target.query)["resource"][-1] + + +def _is_token_endpoint(request: Request) -> bool: + return _identity_resource(request) is not None or request.target.endswith(_EXCHANGE_SUFFIX) + + +def _managed_identity( + resource: str, + hold: _Hold | None, + lifetime_seconds: int = _TOKEN_LIFETIME_SECONDS, + access_token: str | None = None, +) -> Reply: + if hold is not None and resource == hold.resource: + hold.arrivals.put(resource) + hold.release.wait(_HOLD_SECONDS) + return _json( + { + "access_token": access_token or _identity_token(resource), + "expires_on": str(int(time.time()) + lifetime_seconds), + "resource": resource, + "token_type": "Bearer", + } + ) + + +def _entra_exchange(request: Request) -> Reply: + assertion: Final = parse_qs(request.body.decode())["client_assertion"][-1] + return _json({"access_token": f"aad:{assertion}", "expires_in": 3600, "token_type": "Bearer"}) + + +def _edge(upstream: Callable[[Request], Reply], hold: _Hold | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + resource: Final = _identity_resource(request) + if resource is not None: + return _managed_identity(resource, hold) + if request.target.endswith(_EXCHANGE_SUFFIX): + return _entra_exchange(request) + return upstream(request) + + return respond + + +def _azure_openai(request: Request) -> Reply: + if urlsplit(request.target).path.endswith("/chat/completions"): + return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", _DEPLOYMENT, "azure answered", stream=False) + return responses_reply(f"resp_{uuid.uuid4().hex}", _DEPLOYMENT, "azure answered", stream=False) + + +def _identity_fetches(requests: tuple[Request, ...], resource: str) -> int: + return sum(1 for request in requests if _identity_resource(request) == resource) + + +def _provider_calls(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if not _is_token_endpoint(request)) + + +def _marker() -> str: + return f"entra audit {uuid.uuid4().hex}" + + +def _model(scenario: Scenario, surface: _Surface, api_base: str, scope: str) -> str: + match surface: + case _Surface.LEGACY_CHAT | _Surface.RESPONSES: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", + api_base=api_base, + api_key=None, + api_version=_DATED_API_VERSION, + azure_scope=scope, + ) + case _Surface.V1_CHAT: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", api_base=api_base, api_key=None, api_version="v1", azure_scope=scope + ) + case _Surface.RESPONSES_WEBSOCKET: + return scenario.model(model=f"azure/{_DEPLOYMENT}", api_base=api_base, api_key=None, azure_scope=scope) + case _Surface.REALTIME | _Surface.REALTIME_HEALTH_CHECK: + return scenario.model( + model_info={"mode": "realtime"}, + model=f"azure/{_DEPLOYMENT}", + api_base=api_base, + api_key=None, + api_version=_REALTIME_API_VERSION, + azure_scope=scope, + ) + + +def _post(surface: _Surface, gateway: Gateway, key: str, model: str) -> httpx.Response: + if surface is _Surface.RESPONSES: + return gateway.request("POST", "/v1/responses", {"model": model, "input": _marker()}, key=key) + return gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": _marker()}]}, key=key + ) + + +async def _realtime_handshake(gateway: Gateway, key: str, model: str, peer: Peer) -> str | None: + proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/realtime?model={model}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=_HOLD_SECONDS, + ): + credentials: Final = await asyncio.to_thread(peer.credentials.get, True, _HOLD_SECONDS) + return credentials.authorization + + +def _upstream_authorization(surface: _Surface, gateway: Gateway, key: str, model: str, wire: Wire, peer: Peer) -> str: + match surface: + case _Surface.LEGACY_CHAT | _Surface.V1_CHAT | _Surface.RESPONSES: + response: Final = _post(surface, gateway, key, model) + assert response.status_code == 200, response.text + (call,) = _provider_calls(wire.drain()) + return call.headers.get("authorization", "") + case _Surface.RESPONSES_WEBSOCKET: + (events,) = asyncio.run(session(str(gateway.client.base_url), key, model, (_marker(),))) + assert completed(events)["status"] == "completed", events + return peer.credentials.get(timeout=_HOLD_SECONDS).authorization or "" + case _Surface.REALTIME: + return asyncio.run(_realtime_handshake(gateway, key, model, peer)) or "" + case _Surface.REALTIME_HEALTH_CHECK: + gateway.request("GET", "/health", params={"model": model}) + return peer.credentials.get(timeout=_HOLD_SECONDS).authorization or "" + + +def _liveliness(gateway: Gateway) -> int | None: + try: + return httpx.get( + f"{str(gateway.client.base_url).rstrip('/')}/health/liveliness", timeout=_LIVELINESS_SECONDS + ).status_code + except httpx.TimeoutException: + return None + + +def _liveliness_while_held(gateway: Gateway, hold: _Hold) -> int | None: + try: + assert hold.arrivals.get(timeout=_HOLD_SECONDS) == hold.resource + return _liveliness(gateway) + finally: + hold.release.set() + + +@pytest.mark.parametrize("surface", tuple(_Surface)) +def test_a_slow_managed_identity_fetch_leaves_the_proxy_answering_other_requests(rig: _Rig, surface: _Surface) -> None: + resource: Final = f"api://audit-hold-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + responses_peer(rig.cert) as peer, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=1) as pool, + ): + api_base: Final = rig.edge if surface in _HTTP_SURFACES else peer.url.removesuffix("/v1") + model: Final = _model(scenario, surface, api_base, f"{resource}/.default") + key: Final = scenario.key() + call: Final = pool.submit(_upstream_authorization, surface, rig.gateway, key, model, wire, peer) + liveliness: Final = _liveliness_while_held(rig.gateway, hold) + authorization: Final = call.result(timeout=_HOLD_SECONDS) + + assert liveliness == 200, f"/health/liveliness did not answer within {_LIVELINESS_SECONDS}s during the fetch" + assert authorization == f"Bearer {_identity_token(resource)}" + + +@pytest.mark.parametrize("surface", (_Surface.LEGACY_CHAT, _Surface.RESPONSES)) +def test_concurrent_cold_requests_share_one_managed_identity_fetch(rig: _Rig, surface: _Surface) -> None: + resource: Final = f"api://audit-herd-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _model(scenario, surface, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + calls: Final = tuple(pool.submit(_post, surface, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + responses: Final = tuple(call.result(timeout=_HOLD_SECONDS) for call in calls) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * _HERD, [r.text for r in responses] + assert _identity_fetches(requests, resource) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer {_identity_token(resource)}" + ] * _HERD + + +def _scripted_identity(resource: str, fetches: tuple[Callable[[], Reply], ...]) -> Callable[[Request], Reply]: + attempts: Final = itertools.count() + edge: Final = _edge(_azure_openai) + + def respond(request: Request) -> Reply: + if _identity_resource(request) != resource: + return edge(request) + return fetches[min(next(attempts), len(fetches) - 1)]() + + return respond + + +def _held_failure(hold: _Hold) -> Reply: + hold.arrivals.put(hold.resource) + hold.release.wait(_HOLD_SECONDS) + return Reply(body=b"not json") + + +def _herd( + rig: _Rig, edge: Callable[[Request], Reply], resource: str, hold: _Hold, warm_up: bool +) -> tuple[tuple[httpx.Response, ...], tuple[Request, ...]]: + with ( + wire_server(edge, port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _model(scenario, _Surface.LEGACY_CHAT, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + warm_up_responses: Final = (_post(_Surface.LEGACY_CHAT, rig.gateway, key, model),) if warm_up else () + calls: Final = tuple(pool.submit(_post, _Surface.LEGACY_CHAT, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + return warm_up_responses + tuple(call.result(timeout=_HOLD_SECONDS) for call in calls), wire.drain() + + +def test_concurrent_requests_inside_the_refresh_window_share_one_refresh_and_keep_the_valid_token(rig: _Rig) -> None: + resource: Final = f"api://audit-refresh-{uuid.uuid4().hex}" + refreshed: Final = f"mi-token-refreshed:{resource}" + hold: Final = _Hold(resource) + edge: Final = _scripted_identity( + resource, + ( + lambda: _managed_identity(resource, None, _INSIDE_REFRESH_WINDOW_SECONDS), + lambda: _managed_identity(resource, hold, access_token=refreshed), + ), + ) + + responses, requests = _herd(rig, edge, resource, hold, warm_up=True) + + assert [response.status_code for response in responses] == [200] * (_HERD + 1), [r.text for r in responses] + assert _identity_fetches(requests, resource) == 2 + assert sorted(call.headers.get("authorization", "") for call in _provider_calls(requests)) == sorted( + [f"Bearer {_identity_token(resource)}"] * _HERD + [f"Bearer {refreshed}"] + ) + + +def test_a_failed_cold_fetch_fails_only_the_request_that_made_it(rig: _Rig) -> None: + resource: Final = f"api://audit-failure-{uuid.uuid4().hex}" + hold: Final = _Hold(resource) + edge: Final = _scripted_identity(resource, (lambda: _held_failure(hold), lambda: _managed_identity(resource, None))) + + responses, _ = _herd(rig, edge, resource, hold, warm_up=False) + + assert sorted(response.status_code == 200 for response in responses) == [False] + [True] * (_HERD - 1), [ + r.text for r in responses + ] + + +def _oidc_model(scenario: Scenario, rig: _Rig, audience: str) -> str: + return scenario.model( + model=f"azure/{_DEPLOYMENT}", + api_base=rig.edge, + api_key=None, + api_version=_DATED_API_VERSION, + azure_ad_token=f"oidc/azure/{audience}", + tenant_id=_TENANT, + client_id=_APP_CLIENT, + ) + + +def _exchanges(requests: tuple[Request, ...]) -> tuple[Request, ...]: + return tuple(request for request in requests if request.target.endswith(_EXCHANGE_SUFFIX)) + + +def test_an_oidc_federated_deployment_fetches_its_assertion_once_across_responses_calls(rig: _Rig) -> None: + audience: Final = f"api://audit-oidc-{uuid.uuid4().hex}" + with wire_server(_edge(_azure_openai), port=rig.port) as wire, rig.gateway.scenario() as scenario: + model: Final = _oidc_model(scenario, rig, audience) + key: Final = scenario.key() + responses: Final = tuple(_post(_Surface.RESPONSES, rig.gateway, key, model) for _ in range(3)) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses] + assert _identity_fetches(requests, audience) == 1 + exchanges: Final = _exchanges(requests) + assert [request.target for request in exchanges] == [f"/{_TENANT}{_EXCHANGE_SUFFIX}"] + assert {name: values[-1] for name, values in parse_qs(exchanges[0].body.decode()).items()} == { + "client_id": _APP_CLIENT, + "grant_type": "client_credentials", + "scope": "https://cognitiveservices.azure.com/.default", + "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", + "client_assertion": _identity_token(audience), + } + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer aad:{_identity_token(audience)}" + ] * 3 + + +def test_concurrent_cold_oidc_requests_share_one_assertion_fetch_and_one_exchange(rig: _Rig) -> None: + audience: Final = f"api://audit-oidc-herd-{uuid.uuid4().hex}" + hold: Final = _Hold(audience) + with ( + wire_server(_edge(_azure_openai, hold), port=rig.port) as wire, + rig.gateway.scenario() as scenario, + ThreadPoolExecutor(max_workers=_HERD) as pool, + ): + model: Final = _oidc_model(scenario, rig, audience) + key: Final = scenario.key() + calls: Final = tuple(pool.submit(_post, _Surface.RESPONSES, rig.gateway, key, model) for _ in range(_HERD)) + try: + eventually( + hold.arrivals.qsize, + lambda arrived: arrived >= _HERD, + seconds=_HERD_WINDOW_SECONDS, + return_last_on_timeout=True, + ) + finally: + hold.release.set() + responses: Final = tuple(call.result(timeout=_HOLD_SECONDS) for call in calls) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * _HERD, [r.text for r in responses] + assert _identity_fetches(requests, audience) == 1 + assert len(_exchanges(requests)) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer aad:{_identity_token(audience)}" + ] * _HERD + + +def _foundry(request: Request) -> Reply: + assert urlsplit(request.target).path == f"{_FOUNDRY_PROJECT}/openai/v1/responses", request.target + text: Final = "Audit result." + return _json( + { + "status": "completed", + "error": None, + "incomplete_details": None, + "output": [ + { + "type": "message", + "content": [ + { + "type": "output_text", + "text": text, + "annotations": [ + { + "type": "url_citation", + "url": _BING_CITATION, + "title": "Bing result", + "start_index": len(text), + "end_index": len(text), + } + ], + } + ], + } + ], + } + ) + + +class _SearchResult(BaseModel): + model_config = ConfigDict(frozen=True) + + url: str + + +class _SearchResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + results: tuple[_SearchResult, ...] + + +def test_bing_grounding_search_mints_its_entra_token_once_across_searches(rig: _Rig) -> None: + with wire_server(_edge(_foundry), port=rig.port) as wire: + responses: Final = tuple( + rig.gateway.request("POST", f"/v1/search/{_BING_TOOL}", {"query": _marker(), "max_results": 1}) + for _ in range(3) + ) + requests: Final = wire.drain() + + assert [response.status_code for response in responses] == [200] * 3, [r.text for r in responses] + assert [ + tuple(result.url for result in _SearchResponse.model_validate_json(response.content).results) + for response in responses + ] == [(_BING_CITATION,)] * 3 + assert _identity_fetches(requests, _BING_RESOURCE) == 1 + assert [call.headers.get("authorization") for call in _provider_calls(requests)] == [ + f"Bearer {_identity_token(_BING_RESOURCE)}" + ] * 3 + + +def _lifecycle(upstream_id: str) -> Callable[[Request], Reply]: + response: Final[dict[str, JsonValue]] = { + "id": upstream_id, + "object": "response", + "created_at": 1, + "model": _DEPLOYMENT, + "output": [], + } + + def respond(request: Request) -> Reply: + path: Final = urlsplit(request.target).path + match request.method: + case "DELETE": + return _json({"id": upstream_id, "object": "response", "deleted": True}) + case "GET" if path.endswith("/input_items"): + return _json({"object": "list", "data": [], "first_id": None, "last_id": None, "has_more": False}) + case "POST" if path.endswith("/cancel"): + return _json({**response, "status": "cancelled"}) + case _: + return _json({**response, "status": "completed"}) + + return respond + + +_LIFECYCLE_ROUTES: Final = ( + ("GET", "/v1/responses/{}", ""), + ("GET", "/v1/responses/{}/input_items", "/input_items"), + ("POST", "/v1/responses/{}/cancel", "/cancel"), + ("DELETE", "/v1/responses/{}", ""), +) + + +def test_every_responses_route_on_an_entra_deployment_reaches_azure_with_the_managed_identity_token( + rig: _Rig, +) -> None: + resource: Final = f"api://audit-routes-{uuid.uuid4().hex}" + upstream_id: Final = f"resp_{uuid.uuid4().hex}" + with wire_server(_edge(_lifecycle(upstream_id)), port=rig.port) as wire, rig.gateway.scenario() as scenario: + model: Final = _model(scenario, _Surface.RESPONSES, rig.edge, f"{resource}/.default") + key: Final = scenario.key() + created: Final = _post(_Surface.RESPONSES, rig.gateway, key, model) + assert created.status_code == 200, created.text + response_id: Final = string_value(JSON_OBJECT.validate_json(created.content)["id"]) + followups: Final = tuple( + rig.gateway.request(method, route.format(response_id), key=key) for method, route, _ in _LIFECYCLE_ROUTES + ) + requests: Final = wire.drain() + + assert [response.status_code for response in followups] == [200] * len(_LIFECYCLE_ROUTES), [ + response.text for response in followups + ] + calls: Final = _provider_calls(requests) + assert (calls[0].method, urlsplit(calls[0].target).path.endswith("/responses")) == ("POST", True) + assert [ + (call.method, urlsplit(call.target).path.endswith(f"/{upstream_id}{suffix}")) + for call, (_, _, suffix) in zip(calls[1:], _LIFECYCLE_ROUTES, strict=True) + ] == [(method, True) for method, _, _ in _LIFECYCLE_ROUTES] + assert {call.headers.get("authorization") for call in calls} == {f"Bearer {_identity_token(resource)}"} + assert _identity_fetches(requests, resource) == 1 diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index a4542f1f604..c392e355708 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1,4 +1,5 @@ import asyncio +import base64 import json import re import threading @@ -32,6 +33,7 @@ from integration.providers._count_tokens_system_lift import ( LIFTED, MID_SYSTEM, REMINDER, + RESPONSES_REQUEST_OPTIONS, RESPONSES_TOOLS, STRING_CASE, USER, @@ -1702,6 +1704,72 @@ def test_responses_input_tokens_sends_mantle_the_tool_use_invoke_inference_sends assert _tool_use_fields(mantle_count) == expected, mantle_count +_RUNTIME_COUNTED_OPTIONS: Final[Mapping[str, dict[str, JsonValue]]] = MappingProxyType( + {**RESPONSES_REQUEST_OPTIONS, "forced-tool-choice": _FORCED_TOOL_CHOICE} +) +_CONVERSE_WEATHER_TOOLS: Final[dict[str, JsonValue]] = { + "tools": [ + { + "toolSpec": { + "name": "get_weather", + "description": "Look up the current weather for a city", + "inputSchema": {"json": WEATHER_SCHEMA}, + } + } + ] +} +_INVOKE_WEATHER_TOOLS: Final[list[JsonValue]] = [ + { + "type": "custom", + "name": "get_weather", + "description": "Look up the current weather for a city", + "input_schema": WEATHER_SCHEMA, + } +] +_RUNTIME_COUNTED_SHAPES: Final[Mapping[str, tuple[frozenset[str], str, JsonValue]]] = MappingProxyType( + { + "converse": (frozenset({"messages", "system", "toolConfig"}), "toolConfig", _CONVERSE_WEATHER_TOOLS), + "invokeModel": ( + frozenset({"anthropic_version", "max_tokens", "messages", "system", "tools"}), + "tools", + _INVOKE_WEATHER_TOOLS, + ), + } +) + + +def _runtime_counted_request(count_body: Mapping[str, JsonValue]) -> tuple[str, dict[str, JsonValue]]: + count_input: Final = _JSON_OBJECT.validate_python(count_body["input"]) + invoke: Final = count_input.get("invokeModel") + if invoke is None: + return "converse", _JSON_OBJECT.validate_python(count_input["converse"]) + encoded: Final = _JSON_OBJECT.validate_python(invoke)["body"] + assert isinstance(encoded, str), invoke + return "invokeModel", _JSON_OBJECT.validate_json(base64.b64decode(encoded)) + + +@pytest.mark.parametrize("request_body", _RUNTIME_COUNTED_OPTIONS.values(), ids=_RUNTIME_COUNTED_OPTIONS.keys()) +def test_responses_input_tokens_sends_the_runtime_the_tools_without_the_options_it_refuses( + counting_proxy: OwnedProxy, mantle_port: int, request_body: dict[str, JsonValue] +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _metadata_deployment(scenario, runtime.url, _SONNET) + response: Final = gateway.request("POST", "/v1/responses/input_tokens", {**request_body, "model": model}) + assert _payload(response)["input_tokens"] == _RUNTIME_COUNT, response.text + ((target, count_body),) = _runtime_count_requests(runtime) + assert target == f"/model/{_SONNET}/count-tokens", target + shape, counted = _runtime_counted_request(count_body) + allowed_fields, tools_field, weather_tools = _RUNTIME_COUNTED_SHAPES[shape] + assert set(counted) <= allowed_fields, counted + assert counted.get(tools_field) == (weather_tools if "tools" in request_body else None), counted + assert mantle.drain() == () + + def _counting_every_model(request: Request) -> Reply: if _COUNT_TARGET.match(target_of(request)): return _json_reply(200, {"inputTokens": _RUNTIME_COUNT}) diff --git a/tests/integration/providers/test_count_tokens_model_wire.py b/tests/integration/providers/test_count_tokens_model_wire.py index b6ace7ddef3..f15545f4374 100644 --- a/tests/integration/providers/test_count_tokens_model_wire.py +++ b/tests/integration/providers/test_count_tokens_model_wire.py @@ -24,7 +24,10 @@ _SONNET: Final = "us.anthropic.claude-sonnet-5-5" _SONNET_BASE: Final = "anthropic.claude-sonnet-5-5" _OPUS: Final = "us.anthropic.claude-opus-5-5" _OPUS_BASE: Final = "anthropic.claude-opus-5-5" +_OPUS_4_7: Final = "us.anthropic.claude-opus-4-7" _NOVA: Final = "amazon.nova-lite-v1:0" +_RUNTIME_ROUTED: Final = "us.anthropic.claude-sonnet-4-6" +_CONFIGURED_COUNT_REGION: Final = "us-east-2" _KIMI: Final = "us.moonshotai.kimi-k3" _KIMI_BASE: Final = "moonshotai.kimi-k3" _PROFILE_ID: Final = "lit9563count" @@ -168,6 +171,7 @@ def _mantle(request: Request) -> Reply: return _json_reply(200, {"input_tokens": _COUNT_MODEL_COUNT}) error_type: Final = { 401: "authentication_error", + 403: "permission_error", 404: "not_found_error", 429: "rate_limit_error", }[status] @@ -206,9 +210,9 @@ def _mantle_count_requests(wire: Wire) -> tuple[tuple[tuple[str, str], dict[str, return tuple((_signer(request), _JSON_OBJECT.validate_json(request.body)) for request in requests) -def _mantle_body(prompt: str) -> dict[str, JsonValue]: +def _mantle_body(prompt: str, model: str = _OPUS_BASE) -> dict[str, JsonValue]: return { - "model": _OPUS_BASE, + "model": model, "messages": [{"role": "user", "content": prompt}], "system": _SYSTEM, "tools": _SDK_TOOLS, @@ -219,13 +223,13 @@ def _url(port: int) -> str: return f"http://127.0.0.1:{port}" -def _litellm_params(model: str, api_base: str) -> dict[str, JsonValue]: +def _litellm_params(model: str, api_base: str, region: str = _REGION) -> dict[str, JsonValue]: return { "model": f"bedrock/{model}", "api_base": api_base, "aws_access_key_id": _ACCESS_KEY, "aws_secret_access_key": _SECRET_KEY, - "aws_region_name": _REGION, + "aws_region_name": region, } @@ -259,9 +263,13 @@ def counting_proxy(tmp_path_factory: pytest.TempPathFactory, mantle_port: int) - def _deployment( - scenario: Scenario, runtime_port: int, model: str = _SONNET, model_info: Mapping[str, JsonValue] | None = None + scenario: Scenario, + runtime_port: int, + model: str = _SONNET, + model_info: Mapping[str, JsonValue] | None = None, + region: str = _REGION, ) -> str: - return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port))) + return scenario.model(model_info=model_info, api_key=None, **_litellm_params(model, _url(runtime_port), region)) def _count(gateway: Gateway, model: str, status: int, count_model_status: int | None = None) -> MessageTokensCount: @@ -382,6 +390,171 @@ def test_framing_on_one_deployment_leaves_a_sibling_of_the_same_model_raising_th assert _mantle_models(mantle) == () +def test_framing_is_added_to_the_local_estimate_when_a_shipped_runtime_count_model_cannot_count( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + assert _shipped_token_count_model(_RUNTIME_ROUTED).get("region") is None + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + framed: Final = _count( + gateway, + _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_framing": _FRAMING}), + 400, + count_model_status=404, + ) + unframed: Final = _count( + gateway, + _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_framing": 0}), + 400, + count_model_status=404, + ) + assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed) + assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed) + assert _runtime_models(runtime) == (count_model, count_model) + assert _mantle_models(mantle) == (count_model, count_model) + + +_OPUS_4_7_PINNED: Final = ( + _OPUS_4_7, + _REGION, + _shipped_token_count_model_field(_OPUS_4_7, "region"), + "anthropic.claude-opus-4-7", +) + + +@pytest.mark.parametrize( + ("model", "deployment_region", "count_region", "count_model", "count_model_status"), + [ + (_SONNET, _REGION, _COUNT_REGION, _OPUS_BASE, 404), + (_SONNET_BASE, "ca-central-1", "ca-central-1", _OPUS_BASE, 404), + (*_OPUS_4_7_PINNED, 404), + (*_OPUS_4_7_PINNED, 401), + (*_OPUS_4_7_PINNED, 403), + ], + ids=[ + "sonnet_5_5_in_its_pinned_region", + "sonnet_5_5_outside_the_pinned_geography", + "opus_4_7_in_its_pinned_region", + "opus_4_7_pinned_region_denies_the_key", + "opus_4_7_pinned_region_denies_the_account", + ], +) +def test_framing_is_added_to_the_local_estimate_when_a_shipped_mantle_count_model_cannot_count( + counting_proxy: OwnedProxy, + runtime_port: int, + mantle_port: int, + model: str, + deployment_region: str, + count_region: str, + count_model: str, + count_model_status: int, +) -> None: + gateway: Final = counting_proxy.gateway + assert _shipped_token_count_model_field(model, "litellm_provider") == "bedrock_mantle" + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + framed: Final = _count( + gateway, + _deployment(scenario, runtime_port, model, {"token_count_framing": _FRAMING}, deployment_region), + 200, + count_model_status=count_model_status, + ) + unframed: Final = _count( + gateway, + _deployment(scenario, runtime_port, model, {"token_count_framing": 0}, deployment_region), + 200, + count_model_status=count_model_status, + ) + assert (framed.model_extra, unframed.model_extra) == (_LOCAL_SOURCE, _LOCAL_SOURCE), (framed, unframed) + assert framed.input_tokens - unframed.input_tokens == _FRAMING, (framed, unframed) + assert _runtime_counts(runtime) == () + assert ( + _mantle_count_requests(mantle) + == (((_ACCESS_KEY, count_region), _mantle_body(_prompt(200, count_model_status), count_model)),) * 2 + ) + + +def test_framing_keeps_the_403_when_the_shipped_count_model_runs_in_the_deployment_region( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + with pytest.raises(anthropic.PermissionDeniedError, match=re.escape(f"scripted mantle 403 for {_OPUS_BASE}")): + _count( + gateway, + _deployment(scenario, runtime_port, _SONNET_BASE, {"token_count_framing": _FRAMING}, "ca-central-1"), + 200, + count_model_status=403, + ) + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == ( + ((_ACCESS_KEY, "ca-central-1"), _mantle_body(_prompt(200, 403))), + ) + + +def test_configured_token_count_model_error_stands_with_framing( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + model_info: Final[dict[str, JsonValue]] = { + "token_count_framing": _FRAMING, + "token_count_model": {"model": count_model}, + } + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + with pytest.raises(anthropic.NotFoundError, match=re.escape(f"scripted mantle 404 for {count_model}")): + _count( + gateway, _deployment(scenario, runtime_port, _RUNTIME_ROUTED, model_info), 400, count_model_status=404 + ) + assert _runtime_models(runtime) == (count_model,) + assert _mantle_models(mantle) == (count_model,) + + +def test_configured_token_count_model_provider_counts_on_mantle_without_trying_the_runtime( + counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + count_model: Final = _shipped_token_count_model_field(_RUNTIME_ROUTED, "model") + assert _shipped_token_count_model_field(_RUNTIME_ROUTED, "litellm_provider") == "bedrock" + configured: Final[dict[str, JsonValue]] = { + "model": count_model, + "litellm_provider": "bedrock_mantle", + "region": _CONFIGURED_COUNT_REGION, + } + with ( + wire_server(_runtime, port=runtime_port) as runtime, + wire_server(_mantle, port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + counted: Final = _count( + gateway, _deployment(scenario, runtime_port, _RUNTIME_ROUTED, {"token_count_model": configured}), 400 + ) + assert (counted.input_tokens, counted.model_extra) == ( + _COUNT_MODEL_COUNT, + {"litellm_count_source": "token_count_model", "litellm_count_model": count_model}, + ), counted + assert _runtime_counts(runtime) == () + assert _mantle_count_requests(mantle) == ( + ((_ACCESS_KEY, _CONFIGURED_COUNT_REGION), {**_mantle_body(_prompt(400)), "model": count_model}), + ) + + def test_empty_configured_token_count_model_counts_the_model_itself_in_the_deployment_region( counting_proxy: OwnedProxy, runtime_port: int, mantle_port: int ) -> None: 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_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 96410b03a92..9ecec3686a9 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -1,5 +1,7 @@ import json +import uuid from collections.abc import Iterator, Mapping +from dataclasses import dataclass from typing import Final import pytest @@ -488,3 +490,103 @@ def test_responses_input_tokens_sends_openai_the_owners_previous_response_id(cou assert counted.status_code == 200, counted.text count_body: Final = _body_for(wire.drain(), "/v1/responses/input_tokens") assert count_body["previous_response_id"] == _RESPONSES_REPLY["id"] + + +_PROVIDER_FILE_ID: Final = "file-audit-provider" +_UPLOADED_FILE: Final[Mapping[str, JsonValue]] = { + "id": _PROVIDER_FILE_ID, + "object": "file", + "bytes": 5, + "created_at": 1, + "filename": "notes.txt", + "purpose": "user_data", + "status": "processed", +} + + +def _files_peer(request: Request) -> Reply: + if request.target == "/v1/files": + return Reply(body=json.dumps(_UPLOADED_FILE).encode()) + return _count_response(request) + + +def _file_input(file_id: str) -> list[JsonValue]: + return [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Summarize the notes."}, + {"type": "input_file", "file_id": file_id}, + ], + } + ] + + +@dataclass(frozen=True, slots=True) +class _FileTenants: + model_id: str + model: str + owner_key: str + outsider_key: str + file_id: str + + +def _file_tenants(scenario: Scenario, wire: Wire) -> _FileTenants: + model_id: Final = f"audit-openai-{uuid.uuid4().hex}" + model: Final = scenario.model( + model_info={"id": model_id}, + model="openai/gpt-4o-mini", + api_base=f"{wire.url}/v1", + api_key="sk-integration-openai", + ) + owner_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), team_id=scenario.team(models=[model]), models=[model] + ) + outsider_key: Final = scenario.key( + user_id=scenario.user(user_role="internal_user"), team_id=scenario.team(models=[model]), models=[model] + ) + uploaded: Final = scenario.gateway.request_multipart( + "/v1/files", + {"purpose": "user_data", "target_model_names": model}, + {"file": ("notes.txt", b"notes", "text/plain")}, + key=owner_key, + ) + assert uploaded.status_code == 200, uploaded.text + file_id: Final = _JSON_OBJECT.validate_json(uploaded.content)["id"] + assert isinstance(file_id, str) and file_id != _PROVIDER_FILE_ID, uploaded.text + assert [request.target for request in wire.drain()] == ["/v1/files"] + return _FileTenants(model_id, model, owner_key, outsider_key, file_id) + + +def test_responses_input_tokens_sends_openai_the_provider_id_of_a_managed_file(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_files_peer) as wire, gateway.scenario() as scenario: + tenants: Final = _file_tenants(scenario, wire) + response: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + { + "model": tenants.model, + "input": _file_input(tenants.file_id), + "model_file_id_mapping": {tenants.file_id: {tenants.model_id: "file-forged"}}, + }, + key=tenants.owner_key, + ) + assert response.status_code == 200, response.text + assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT + assert _count_body(wire) == {"model": "gpt-4o-mini", "input": _file_input(_PROVIDER_FILE_ID)} + + +def test_responses_input_tokens_refuses_another_teams_managed_file_before_openai(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_files_peer) as wire, gateway.scenario() as scenario: + tenants: Final = _file_tenants(scenario, wire) + refused: Final = gateway.request( + "POST", + "/v1/responses/input_tokens", + {"model": tenants.model, "input": _file_input(tenants.file_id)}, + key=tenants.outsider_key, + ) + assert refused.status_code == 403, refused.text + assert _JSON_OBJECT.validate_json(refused.content)["error"]["type"] == "permission_error", refused.text + assert wire.drain() == () diff --git a/tests/integration/routing/test_end_user_region_routing.py b/tests/integration/routing/test_end_user_region_routing.py index ce90280f173..884b61228cc 100644 --- a/tests/integration/routing/test_end_user_region_routing.py +++ b/tests/integration/routing/test_end_user_region_routing.py @@ -19,6 +19,9 @@ UNCOUNTABLE_MODEL: Final = "regional-uncountable-model" PINNED_REGION: Final = "eu" WEIGHT_BY_REGION: Final = {"eu": 1, "us": 99} SILENT_WAV: Final = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQAAAAA=" +RPM_CAPPED_MODEL: Final = "regional-rpm-capped-model" +RPM_CAPPED_DEPLOYMENTS: Final = frozenset(f"{RPM_CAPPED_MODEL}-{region}" for region in UPSTREAM_BY_REGION) +RPM_CAPPED_CALLS: Final = 3 @pytest.fixture(scope="module") @@ -53,6 +56,19 @@ def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: } for region, upstream in UPSTREAM_BY_REGION.items() ), + *( + { + "model_name": RPM_CAPPED_MODEL, + "litellm_params": { + "model": f"openai/{upstream}", + "api_base": f"{base.upstream_url}/v1", + "api_key": "synthetic-region-key", + "rpm": 1, + }, + "model_info": {"id": f"{RPM_CAPPED_MODEL}-{region}"}, + } + for region, upstream in UPSTREAM_BY_REGION.items() + ), ] config["router_settings"] = {**config["router_settings"], "enable_pre_call_checks": True} path: Final = directory / "region-routing.yaml" @@ -143,3 +159,16 @@ def test_an_uncountable_prompt_still_pins_an_end_user_to_their_region( observed: Final = upstream.get("/__observations").json()["requests"] pinned_upstream: Final = UPSTREAM_BY_REGION[PINNED_REGION] assert [request["body"]["model"] for request in observed] == [pinned_upstream] * recorded_chat_calls + + +def test_pre_call_checks_keep_serving_deployments_past_their_rpm(candidate: Gateway) -> None: + with candidate.scenario() as scenario: + key: Final = scenario.key(models=[RPM_CAPPED_MODEL]) + body: Final = {"model": RPM_CAPPED_MODEL, "messages": [{"role": "user", "content": "integration control"}]} + responses: Final = tuple( + candidate.request("POST", "/v1/chat/completions", body, key=key) for _ in range(RPM_CAPPED_CALLS) + ) + assert [response.status_code for response in responses] == [200] * RPM_CAPPED_CALLS, [ + response.text for response in responses + ] + assert {response.headers.get("x-litellm-model-id") for response in responses} <= RPM_CAPPED_DEPLOYMENTS 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/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/test_batch_poll_starvation.py b/tests/integration/spend/test_batch_poll_starvation.py index 51f5f0871f5..f93cb807907 100644 --- a/tests/integration/spend/test_batch_poll_starvation.py +++ b/tests/integration/spend/test_batch_poll_starvation.py @@ -365,8 +365,10 @@ async def test_arrivals_progress_while_retry_scan_advances_and_wraps(count: int) assert any(latest_ids[0] in tuple(row.id for row in page) for page in followups) restarted: Final = CheckBatchCost(logging, prisma, router) - restart_page: Final = await restarted._find_jobs(where) - assert tuple(row.id for row in restart_page) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE] + restart_ids: Final = tuple(row.id for row in await restarted._find_jobs(where)) + restart_old: Final = tuple(identity for identity in restart_ids if identity in old_ids) + assert len(restart_ids) == MAX_OBJECTS_PER_POLL_CYCLE + assert restart_old == old_ids[: MAX_OBJECTS_PER_POLL_CYCLE - MAX_OBJECTS_PER_POLL_CYCLE // 2] finally: write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,)) 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/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/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_file_deletion_blocking.py b/tests/unit/enterprise/proxy/test_file_deletion_blocking.py index b7acdad35a1..4ca5e645f7f 100644 --- a/tests/unit/enterprise/proxy/test_file_deletion_blocking.py +++ b/tests/unit/enterprise/proxy/test_file_deletion_blocking.py @@ -10,6 +10,8 @@ This ensures cost tracking is not disrupted by premature file deletion. import base64 import json +from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -17,6 +19,8 @@ from fastapi import HTTPException from litellm.proxy._types import UserAPIKeyAuth +NOW: Final = datetime(2026, 10, 1, tzinfo=timezone.utc) + def _make_unified_file_id(file_id: str = "file-abc123") -> str: """Create a base64-encoded unified file ID.""" @@ -50,7 +54,7 @@ def _make_batch_db_record( mock_batch.status = status mock_batch.file_object = json.dumps(file_object) mock_batch.created_by = created_by - mock_batch.created_at = 1700000000 + mock_batch.created_at = NOW return mock_batch @@ -105,6 +109,7 @@ def _make_managed_files_instance_with_batches( instance = PROXY_LiteLLMManagedFiles( internal_usage_cache=mock_cache, prisma_client=mock_prisma, + clock=lambda: NOW, ) return instance 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/images/test_image_edit.py b/tests/unit/images/test_image_edit.py index d8f1990af0e..fb7a2869b2a 100644 --- a/tests/unit/images/test_image_edit.py +++ b/tests/unit/images/test_image_edit.py @@ -12,7 +12,7 @@ from typing_extensions import ReadOnly, TypedDict, override import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.types.utils import ImageResponse +from litellm.types.utils import FileTypes, ImageResponse _PNG_SIGNATURE: Final = b"\x89PNG\r\n\x1a\n" _FIRST_IMAGE: Final = _PNG_SIGNATURE + b"first-reference-image" @@ -73,6 +73,20 @@ def _multipart_image_parts(request: httpx.Request) -> tuple[bytes, ...]: return tuple(part.split(b"\r\n\r\n", 1)[1].removesuffix(b"\r\n") for part in parts if b'name="image[]"' in part) +@pytest.mark.asyncio +@pytest.mark.parametrize("image", (None, []), ids=("none", "empty-list")) +async def test_bfl_image_edit_missing_image_returns_bad_request(image: list[FileTypes] | None) -> None: + with pytest.raises(litellm.BadRequestError, match="No image provided") as exc_info: + await litellm.aimage_edit( + model="black_forest_labs/flux-kontext-pro", + prompt="make it blue", + image=image, + api_key="fake-key", + ) + + assert exc_info.value.status_code == 400 + + def test_openai_image_edit_sync_accepts_bytesio_images(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) 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_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/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..a1b58521d7e 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, @@ -1367,6 +1370,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 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/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/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/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index 63f149fabe4..89352ce8702 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -569,6 +569,7 @@ _ANTHROPIC_TOOL_USE: Final[list[dict[str, JsonValue]]] = [ ("messages", "tool", "messages_format"), ( (_ANTHROPIC_TOOL_USE, {"name": "get_weather", "input_schema": _WEATHER_SCHEMA}, "openai"), + ([{"role": "user", "content": "Open example.com"}], {"type": "computer_toolset_20260801"}, "openai"), ( [{"role": "user", "content": "Weather in SF?"}], {"type": "function", "name": "get_weather", "parameters": _WEATHER_SCHEMA}, @@ -685,6 +686,26 @@ _CHAT_WEATHER_TOOL: Final[dict[str, JsonValue]] = { [{"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}, _CHAT_WEATHER_TOOL], id="mcp-tool", ), + pytest.param( + [{"role": "user", "content": "Open the docs"}], + [{"type": "computer_use_preview", "display_width": 1024, "display_height": 768, "environment": "browser"}], + id="computer-use-preview-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Ask the wiki"}], + [{"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}], + id="mcp-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Plot the sales data"}], + [{"type": "code_interpreter", "container": {"type": "auto"}}], + id="code-interpreter-tool-alone", + ), + pytest.param( + [{"role": "user", "content": "Fix the failing test"}], + [{"type": "apply_patch"}], + id="apply-patch-tool-alone", + ), ), ) async def test_counter_rejects_input_chat_translation_would_leave_out_without_calling_anthropic( @@ -700,6 +721,80 @@ async def test_counter_rejects_input_chat_translation_would_leave_out_without_ca assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +async def test_counter_sends_a_tool_with_a_non_string_type_to_anthropic_as_given( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], count_endpoint: str +) -> None: + tool: Final[dict[str, JsonValue]] = {"type": ["mcp"], "name": "lookup"} + + result, bodies = await _counted_body( + counter_type, count_endpoint, [{"role": "user", "content": "Ask the wiki"}], [tool], "openai" + ) + + assert result is not None + assert [body.get("tools") for body in bodies] == [[tool]] + + +_CRM_NAMESPACE_TOOL: Final[dict[str, JsonValue]] = { + "type": "namespace", + "name": "crm", + "description": "CRM lookups", + "tools": [ + { + "type": "function", + "name": "find_owner", + "description": "Find the owner of an account", + "parameters": {"type": "object", "properties": {"account": {"type": "string"}}, "required": ["account"]}, + } + ], +} + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport_clients") +@_COUNTERS +@pytest.mark.parametrize( + "tools", + ( + pytest.param([{"type": "tool_search"}], id="tool-search-alone"), + pytest.param( + [ + {"type": "tool_search"}, + { + "type": "function", + "function": {"name": "get_weather", "parameters": _WEATHER_SCHEMA}, + "defer_loading": True, + }, + ], + id="tool-search-with-deferred-function", + ), + pytest.param([_CRM_NAMESPACE_TOOL], id="namespace-alone"), + pytest.param([_CHAT_WEATHER_TOOL, _CRM_NAMESPACE_TOOL], id="namespace-with-function"), + ), +) +async def test_counter_sends_the_tools_chat_inference_makes_from_tool_search_and_namespace( + counter_type: type[AnthropicTokenCounter] | type[AzureAIAnthropicTokenCounter], + count_endpoint: str, + tools: list[dict[str, JsonValue]], +) -> None: + messages: Final[list[dict[str, JsonValue]]] = [{"role": "user", "content": "Who owns the north region account?"}] + config: Final = AnthropicConfig() + inference: Final = config.transform_request( + model="claude-opus-5-5", + messages=deepcopy(messages), + optional_params=config.map_openai_params({"tools": deepcopy(tools)}, {}, "claude-opus-5-5", drop_params=False), + litellm_params={}, + headers={}, + ) + + result, bodies = await _counted_body(counter_type, count_endpoint, messages, tools, "openai") + + assert result is not None + assert [body.get("tools") for body in bodies] == [inference["tools"]] + + _ANSWER_SCHEMA: Final[dict[str, JsonValue]] = { "type": "object", "properties": {"answer": {"type": "string"}}, diff --git a/tests/unit/llms/azure/test_azure.py b/tests/unit/llms/azure/test_azure.py index 7d5afde68ef..c52321baf02 100644 --- a/tests/unit/llms/azure/test_azure.py +++ b/tests/unit/llms/azure/test_azure.py @@ -6,7 +6,7 @@ import os import time from collections.abc import Mapping from typing import Final -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import httpx import pytest @@ -291,11 +291,9 @@ def test_get_azure_ad_token_from_username_password(mock_get_bearer_token_provide mock_credential.assert_called_once_with(client_id=client_id, username=username, password=password) - mock_get_bearer_token_provider.assert_called_once_with( - mock_credential.return_value, "https://cognitiveservices.azure.com/.default" - ) + mock_get_bearer_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default") - assert result == mock_token_provider + assert result() == "mock-token" def test_azure_openai_gpt_4o_naming(monkeypatch): diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index 1dc7ff74323..f7e2e73d426 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -2,9 +2,12 @@ import json import os import threading import traceback +import uuid +from concurrent.futures import ThreadPoolExecutor +from queue import SimpleQueue from typing import Callable, Final, Optional from urllib.parse import parse_qs -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import httpx import pytest @@ -13,6 +16,7 @@ import litellm from litellm.llms.azure.common_utils import ( BaseAzureLLM, _cached_azure_ad_token_refresh_provider, + _cached_default_azure_credential_provider, _cached_entra_id_token_provider, get_azure_ad_token, get_azure_ad_token_from_entra_id, @@ -41,6 +45,7 @@ def setup_mocks(monkeypatch): monkeypatch.delenv("AZURE_SCOPE", raising=False) monkeypatch.delenv("AZURE_AD_TOKEN", raising=False) _cached_azure_ad_token_refresh_provider.cache_clear() + _cached_default_azure_credential_provider.cache_clear() with ( patch( @@ -86,6 +91,7 @@ def setup_mocks(monkeypatch): "select_url": mock_select_url, } _cached_azure_ad_token_refresh_provider.cache_clear() + _cached_default_azure_credential_provider.cache_clear() def test_initialize_with_api_key(setup_mocks): @@ -1436,6 +1442,33 @@ def test_get_azure_ad_token_with_token_refresh_reuses_provider_per_scope(setup_m ] +def test_get_azure_ad_token_reuses_one_default_azure_credential_provider_per_scope(setup_mocks): + setup_mocks["litellm"].enable_azure_ad_token_refresh = True + default_scope: Final = "https://cognitiveservices.azure.com/.default" + other_scope: Final = "api://other-resource/.default" + + def provider_for(azure_scope: str, azure_credential: AzureCredentialType | None = None) -> Callable[[], str]: + if azure_credential != AzureCredentialType.DefaultAzureCredential: + raise ValueError("No credential provided") + return lambda: f"token-for-{azure_scope}" + + setup_mocks["token_provider"].side_effect = provider_for + + tokens: Final = ( + get_azure_ad_token(GenericLiteLLMParams()), + get_azure_ad_token(GenericLiteLLMParams()), + get_azure_ad_token(GenericLiteLLMParams(azure_scope=other_scope)), + ) + + default_credential_builds: Final = tuple( + call.kwargs["azure_scope"] + for call in setup_mocks["token_provider"].call_args_list + if call.kwargs.get("azure_credential") == AzureCredentialType.DefaultAzureCredential + ) + assert tokens == (f"token-for-{default_scope}", f"token-for-{default_scope}", f"token-for-{other_scope}") + assert default_credential_builds == (default_scope, other_scope) + + def test_get_azure_ad_token_with_token_refresh_error(setup_mocks): """Test get_azure_ad_token with token refresh enabled but raising an error.""" # Reset mocks to ensure clean state @@ -1536,9 +1569,7 @@ def test_get_azure_ad_token_provider_with_default_azure_credential(): mock_default_cred.assert_called_once_with() # Verify get_bearer_token_provider was called with the right parameters - mock_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" - ) + mock_token_provider.assert_called_once_with(ANY, "https://cognitiveservices.azure.com/.default") # Verify the returned token provider works token = token_provider() @@ -2137,6 +2168,42 @@ class TestEntraIdTokenProviderCache: assert first is second assert mock_credential.call_count == 1 + def test_concurrent_first_callers_build_one_credential(self): + callers: Final = 8 + start: Final = threading.Barrier(callers) + built: Final = SimpleQueue[None]() + everyone_built: Final = threading.Event() + release: Final = threading.Event() + + def held_credential(*_: str) -> MagicMock: + built.put(None) + if built.qsize() >= callers: + everyone_built.set() + release.wait(5) + return MagicMock() + + def first_call() -> Callable[[], str]: + start.wait(5) + return get_azure_ad_token_from_entra_id( + tenant_id="tenant", + client_id="client", + client_secret="secret", + scope="https://cognitiveservices.azure.com/.default", + ) + + with ( + patch("azure.identity.ClientSecretCredential", side_effect=held_credential), + patch("azure.identity.get_bearer_token_provider", side_effect=lambda credential, scope: lambda: "token"), + ThreadPoolExecutor(max_workers=callers) as pool, + ): + calls: Final = tuple(pool.submit(first_call) for _ in range(callers)) + everyone_built.wait(0.5) + release.set() + providers: Final = tuple(call.result(timeout=5) for call in calls) + + assert built.qsize() == 1 + assert all(provider is providers[0] for provider in providers) + @pytest.mark.parametrize( "second_call_kwargs", [ @@ -2291,6 +2358,71 @@ def test_oidc_token_exchange_closes_its_own_http_client(monkeypatch: pytest.Monk assert token_endpoint.closed +def test_concurrent_oidc_callers_share_one_token_exchange(monkeypatch: pytest.MonkeyPatch) -> None: + callers: Final = 8 + start: Final = threading.Barrier(callers) + exchanges: Final = SimpleQueue[None]() + everyone_exchanging: Final = threading.Event() + release: Final = threading.Event() + + class HeldTokenEndpoint(httpx.BaseTransport): + def handle_request(self, request: httpx.Request) -> httpx.Response: + exchanges.put(None) + if exchanges.qsize() >= callers: + everyone_exchanging.set() + release.wait(5) + return httpx.Response(200, json={"access_token": "exchanged-token", "expires_in": 3600}) + + monkeypatch.setattr(litellm, "sync_transport", HeldTokenEndpoint(), raising=False) + monkeypatch.setenv("AZURE_TEST_OIDC_BURST_ASSERTION", "burst-jwt") + tenant: Final = f"tenant-{uuid.uuid4().hex}" + + def exchange() -> str: + start.wait(5) + return get_azure_ad_token_from_oidc( + azure_ad_token="os.environ/AZURE_TEST_OIDC_BURST_ASSERTION", + azure_client_id="client", + azure_tenant_id=tenant, + ) + + with ThreadPoolExecutor(max_workers=callers) as pool: + calls: Final = tuple(pool.submit(exchange) for _ in range(callers)) + exchanges.get(timeout=5) + everyone_exchanging.wait(0.5) + release.set() + tokens: Final = tuple(call.result(timeout=5) for call in calls) + + assert tokens == ("exchanged-token",) * callers + assert exchanges.qsize() == 0 + + +def test_oidc_tokens_for_different_scopes_are_exchanged_and_cached_separately(monkeypatch: pytest.MonkeyPatch) -> None: + exchanged_scopes: Final = SimpleQueue[str]() + + class ScopedTokenEndpoint(httpx.BaseTransport): + def handle_request(self, request: httpx.Request) -> httpx.Response: + scope: Final = parse_qs(request.read().decode())["scope"][-1] + exchanged_scopes.put(scope) + return httpx.Response(200, json={"access_token": f"token-for:{scope}", "expires_in": 3600}) + + monkeypatch.setattr(litellm, "sync_transport", ScopedTokenEndpoint(), raising=False) + monkeypatch.setenv("AZURE_TEST_OIDC_SCOPED_ASSERTION", "scoped-jwt") + tenant: Final = f"tenant-{uuid.uuid4().hex}" + + def token_for(scope: str) -> str: + return get_azure_ad_token_from_oidc( + azure_ad_token="os.environ/AZURE_TEST_OIDC_SCOPED_ASSERTION", + azure_client_id="client", + azure_tenant_id=tenant, + scope=scope, + ) + + tokens: Final = tuple(token_for(scope) for scope in ("api://first/.default", "api://second/.default") * 2) + + assert tokens == ("token-for:api://first/.default", "token-for:api://second/.default") * 2 + assert exchanged_scopes.qsize() == 2 + + def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch): """`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client. diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index 6fb47491323..cec7950b859 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -6,6 +6,7 @@ import httpx import pytest from pydantic import JsonValue, TypeAdapter +import litellm from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -264,6 +265,79 @@ async def test_claude_chat_tool_use_is_counted_in_the_shape_claude_inference_sen {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"}]}, ] + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool", + ( + {"type": "mcp", "server_label": "wiki", "server_url": "https://mcp.example/mcp"}, + {"type": "computer_use_preview", "display_width": 1024, "display_height": 768, "environment": "browser"}, + ), + ids=("mcp", "computer_use_preview"), +) +async def test_claude_count_rejects_a_lone_responses_tool_without_calling_bedrock(tool: dict[str, JsonValue]) -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=[{"role": "user", "content": "Ask the wiki"}], + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + tools=[tool], + messages_format="openai", + ) + + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + assert bedrock.posted_hosts() == () + + +_CLAUDE_REPLY: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 3, "output_tokens": 1}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "deployment_model", + ("bedrock/invoke/us.anthropic.claude-sonnet-5-5", "bedrock_mantle/us.anthropic.claude-sonnet-5-5"), +) +async def test_claude_count_sends_the_tool_search_tool_the_deployments_inference_sends(deployment_model: str) -> None: + credentials: Final = DEPLOYMENT["litellm_params"] + inference: Final = _Bedrock(runtime=(200, _CLAUDE_REPLY), mantle=(200, _CLAUDE_REPLY)) + await litellm.acompletion( + model=deployment_model, + messages=MESSAGES, + tools=[{"type": "tool_search"}], + max_tokens=8, + client=AsyncHTTPHandler(transport=httpx.MockTransport(inference)), + **credentials, + ) + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/us.anthropic.claude-opus-5-5", + messages=MESSAGES, + contents=None, + deployment={"litellm_params": {**credentials, "model": deployment_model}}, + request_model="claude-sonnet-5-5", + tools=[{"type": "tool_search"}], + messages_format="openai", + ) + + assert result is not None + assert (result.error, result.total_tokens) == (False, 2177) + counted_tools: Final = _JSON_BODY.validate_json(bedrock.requests[0].content)["tools"] + assert counted_tools == _JSON_BODY.validate_json(inference.requests[0].content)["tools"] + + @pytest.mark.asyncio async def test_claude_chat_tool_use_counted_on_mantle_metadata_route_is_translated_too() -> None: bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) 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/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/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/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index d7fc844171d..5ff159f80d3 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -9,22 +9,27 @@ ARN unified_object_id) batches with no managed unified id. import asyncio import json import logging +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from copy import copy -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from hashlib import sha256 from queue import SimpleQueue from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Literal from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest import respx from fastapi import HTTPException +import litellm +from litellm.constants import MANAGED_OBJECT_STALENESS_CUTOFF_DAYS from litellm.types.router import LiteLLM_Params from litellm.types.utils import LiteLLMBatch +from litellm_enterprise.proxy.common_utils.check_batch_cost import GIVE_UP_POLLING_AFTER if TYPE_CHECKING: from litellm.batches.batch_utils import BatchCostUsageResult @@ -32,6 +37,11 @@ if TYPE_CHECKING: _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" _CLAIM_UNIFIED_BATCH_ID = "dW5pZmllZF9iYXRjaF9pZA==" _CLAIM_OUTPUT_FILE_ID = "file-output-123" +_NOW: Final = datetime(2026, 1, 15, 12, tzinfo=timezone.utc) + + +def _fixed_clock() -> datetime: + return _NOW def test_batch_output_file_object_derives_metadata(): @@ -152,25 +162,9 @@ class TestCheckBatchCost: proxy_logging_obj=mock_proxy_logging_obj, prisma_client=mock_prisma_client, llm_router=mock_llm_router, + clock=_fixed_clock, ) - @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): - """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) - # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) - - await check_batch_cost_instance.check_batch_cost() - - calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list - stale_call = calls[0] - assert stale_call[1]["data"] == {"status": "stale_expired"} - where = stale_call[1]["where"] - assert where["file_purpose"] == "batch" - assert "stale_expired" in where["status"]["not_in"] - assert "created_at" in where - @pytest.mark.asyncio async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) @@ -251,7 +245,7 @@ class TestCheckBatchCost: fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where assert "stale_expired" in fallback_where["status"]["not_in"] - assert calls[1][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE + assert calls[2][1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE # Column absence is now cached — next call should go straight to fallback assert check_batch_cost_instance._has_batch_processed_column is False assert check_batch_cost_instance.batch_processed_support_confirmed is False @@ -744,7 +738,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( "Expected update() to be called exactly once for the completed job" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -877,7 +871,7 @@ class TestCheckBatchCost: mock_update_database.assert_awaited_once() assert mock_update_database.call_args.kwargs["response_cost"] == 0.01 - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( "the job must still be marked processed once cost tracking succeeds" ) @@ -979,7 +973,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( f"Expected update() to be called exactly once for a {terminal_status} job" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -1085,13 +1079,9 @@ class TestCheckBatchCost: assert owner.user_id == "user-1" assert owner.team_id == "team-1" - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1 update_call = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args - assert update_call.kwargs["where"] == { - "id": "job-terminal-mint-1", - "batch_processed": False, - "updated_at": mock_job.updated_at, - } + assert update_call.kwargs["where"] == {"id": "job-terminal-mint-1", "batch_processed": False} update_data = update_call.kwargs["data"] assert update_data["status"] == terminal_status assert update_data["batch_processed"] is True @@ -1156,7 +1146,7 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( "a completed batch with no output file must be marked processed exactly once" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -1205,6 +1195,7 @@ class TestCheckBatchCost: b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" ).decode() mock_job.created_by = "user-1" + mock_job.created_at = _NOW assert check_batch_cost_instance._has_batch_processed_column is True mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) @@ -1245,6 +1236,7 @@ class TestCheckBatchCost: mock_job.id = "job-in-progress-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" + mock_job.created_at = _NOW mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) @@ -1375,7 +1367,7 @@ class TestCheckBatchCost: f"{terminal_status} batch with an output file must fetch results and be billed" ) mock_logging_obj.async_success_handler.assert_awaited_once() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2 + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1 update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == terminal_status, ( @@ -1548,7 +1540,7 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 2, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count == 1, ( "a terminal batch with a 404ing output file must be retired, not retried forever" ) update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] @@ -1946,7 +1938,7 @@ class TestUnmanagedVertexRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01 - assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 1 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2163,7 +2155,7 @@ class TestUnmanagedBedrockRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02 - assert prisma.db.litellm_managedobjecttable.update_many.call_count == 2 + assert prisma.db.litellm_managedobjecttable.update_many.call_count == 1 update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2710,6 +2702,7 @@ class TestPollPageStarvation: proxy_logging_obj=proxy_logging_obj, prisma_client=prisma, llm_router=llm_router, + clock=_fixed_clock, ) def _prisma(self, jobs): @@ -2725,6 +2718,7 @@ class TestPollPageStarvation: job.id = job_id job.unified_object_id = unified_object_id job.created_by = "user-1" + job.created_at = _NOW return job @staticmethod @@ -2910,6 +2904,173 @@ class TestPollPageStarvation: prisma.db.litellm_managedobjecttable.update.assert_not_awaited() + @pytest.mark.asyncio + async def test_a_new_batch_is_polled_every_cycle_while_an_older_batch_stays_unbilled(self) -> None: + """An unbilled batch older than the arrival window used to take every other cycle away from the arrival + lane, so a new batch waited two poll intervals to be billed.""" + now: Final = _NOW + rows: Final = tuple( + SimpleNamespace( + id=f"job-{name}", + unified_object_id=self._encode(f"litellm_proxy;model_id:model-123;llm_batch_id:batch_{name}"), + created_by="user-1", + created_at=now - age, + updated_at=now, + status="in_progress", + ) + for name, age in (("old", timedelta(days=3)), ("new", timedelta(hours=1))) + ) + prisma: Final = self._prisma([]) + prisma.db.litellm_managedobjecttable = _FakeManagedObjectPages(rows) + llm_router: Final = MagicMock() + llm_router.aretrieve_batch = AsyncMock(return_value=MagicMock(status="in_progress")) + poller: Final = self._instance(prisma, llm_router) + + async def _batches_polled_in_one_cycle() -> tuple[str, ...]: + llm_router.aretrieve_batch.reset_mock() + await poller.check_batch_cost() + return tuple(sorted(call.kwargs["batch_id"] for call in llm_router.aretrieve_batch.await_args_list)) + + assert [await _batches_polled_in_one_cycle() for _ in range(4)] == [("batch_new", "batch_old")] * 4 + + def _unbilled_rows(self, ages: dict[str, timedelta]) -> tuple[SimpleNamespace, ...]: + now: Final = _NOW + return tuple( + SimpleNamespace( + id=f"job-{name}", + unified_object_id=self._encode(f"litellm_proxy;model_id:model-123;llm_batch_id:batch_{name}"), + created_by="user-1", + created_at=now - age, + updated_at=now, + status="in_progress", + ) + for name, age in ages.items() + ) + + async def _poll_cycles( + self, monkeypatch: pytest.MonkeyPatch, ages: dict[str, timedelta], cycles: int, page_size: int = 2 + ) -> list[frozenset[str]]: + from litellm_enterprise.proxy.common_utils import check_batch_cost + + monkeypatch.setattr(check_batch_cost, "MAX_OBJECTS_PER_POLL_CYCLE", page_size) + prisma: Final = self._prisma([]) + prisma.db.litellm_managedobjecttable = _FakeManagedObjectPages(self._unbilled_rows(ages)) + llm_router: Final = MagicMock() + llm_router.aretrieve_batch = AsyncMock(return_value=MagicMock(status="in_progress")) + poller: Final = self._instance(prisma, llm_router) + + async def _batches_polled_in_one_cycle() -> frozenset[str]: + llm_router.aretrieve_batch.reset_mock() + await poller.check_batch_cost() + return frozenset(call.kwargs["batch_id"] for call in llm_router.aretrieve_batch.await_args_list) + + return [await _batches_polled_in_one_cycle() for _ in range(cycles)] + + @pytest.mark.asyncio + async def test_an_old_batch_is_polled_every_cycle_while_new_batches_fill_the_page( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + new_batches: Final = {f"new{n}": timedelta(hours=5 - n) for n in range(1, 5)} + polled: Final = await self._poll_cycles(monkeypatch, {"old": timedelta(days=3), **new_batches}, cycles=8) + + assert all("batch_old" in cycle for cycle in polled) + assert all(f"batch_{name}" in frozenset().union(*polled[:4]) for name in new_batches) + + @pytest.mark.asyncio + async def test_old_batches_that_fill_the_page_never_push_a_new_batch_out( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + polled: Final = await self._poll_cycles( + monkeypatch, + {"old1": timedelta(days=3, hours=1), "old2": timedelta(days=3), "new": timedelta(hours=1)}, + cycles=6, + ) + + assert all("batch_new" in cycle for cycle in polled) + assert frozenset({"batch_old1", "batch_old2"}) <= frozenset().union(*polled[:2]) + + @pytest.mark.asyncio + async def test_old_and_new_batches_each_get_half_the_page_when_both_overflow_it( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + ages: Final = { + **{f"old{n}": timedelta(days=3, hours=n) for n in range(3)}, + **{f"new{n}": timedelta(hours=n + 1) for n in range(3)}, + } + polled: Final = await self._poll_cycles(monkeypatch, ages, cycles=6) + + assert [sorted(batch_id[len("batch_") : -1] for batch_id in cycle) for cycle in polled] == [["new", "old"]] * 6 + assert frozenset().union(*polled) == frozenset(f"batch_{name}" for name in ages) + + @pytest.mark.asyncio + async def test_old_and_new_batches_take_turns_on_a_one_row_page(self, monkeypatch: pytest.MonkeyPatch) -> None: + polled: Final = await self._poll_cycles( + monkeypatch, {"old": timedelta(days=3), "new": timedelta(hours=1)}, cycles=4, page_size=1 + ) + + assert polled == [frozenset({"batch_old"}), frozenset({"batch_new"})] * 2 + + @pytest.mark.asyncio + async def test_new_batches_fill_the_page_when_no_old_batch_needs_polling( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + polled: Final = await self._poll_cycles( + monkeypatch, {f"new{n}": timedelta(hours=n + 1) for n in range(3)}, cycles=4 + ) + + assert [len(cycle) for cycle in polled] == [2] * 4 + + @pytest.mark.asyncio + async def test_a_failed_expiry_write_leaves_the_rest_of_the_page_polled(self) -> None: + rows: Final = self._unbilled_rows({"old": _PAST_THE_GIVE_UP_AGE, "new": timedelta(hours=1)}) + table: Final = _FakeManagedObjectPages(rows) + table.update_many = AsyncMock(side_effect=RuntimeError("database unavailable")) + prisma: Final = self._prisma([]) + prisma.db.litellm_managedobjecttable = table + llm_router: Final = MagicMock() + llm_router.aretrieve_batch = AsyncMock(side_effect=_A_REJECTED_KEY) + poller: Final = self._instance(prisma, llm_router) + + await poller.check_batch_cost() + + assert table.update_many.await_count == 1 + assert {call.kwargs["batch_id"] for call in llm_router.aretrieve_batch.await_args_list} == { + "batch_old", + "batch_new", + } + + +class _FakeManagedObjectPages: + """A LiteLLM_ManagedObjectTable double over several unbilled rows that pages them the way Postgres would: + inside the lane's created_at window, after the keyset cursor, in created_at then id order, take at most.""" + + def __init__(self, rows: tuple[SimpleNamespace, ...]) -> None: + self.rows = rows + self.find_many = AsyncMock(side_effect=self._find_many) + self.find_first = AsyncMock(return_value=None) + self.update_many = AsyncMock(return_value=0) + self.update = AsyncMock() + + @staticmethod + def _selected(row: SimpleNamespace, where: Mapping[str, object]) -> bool: + window: Final = where["created_at"] + cursor: Final = where.get("OR") + return ( + row.created_at <= window["lte"] + and ("gt" not in window or row.created_at > window["gt"]) + and ( + cursor is None + or row.created_at > cursor[0]["created_at"]["gt"] + or (row.created_at == cursor[1]["created_at"] and row.id > cursor[1]["id"]["gt"]) + ) + ) + + async def _find_many( + self, *, where: Mapping[str, object], take: int, order: Sequence[Mapping[str, str]] + ) -> tuple[SimpleNamespace, ...]: + selected: Final = (row for row in self.rows if self._selected(row, where)) + return tuple(sorted(selected, key=lambda row: (row.created_at, row.id))[:take]) + class _FakeManagedObjectRow: """One managed batch row the provider has finished but nothing has costed yet.""" @@ -2940,9 +3101,12 @@ class _FakeManagedObjectRow: class _FakeManagedObjectTable: """A LiteLLM_ManagedObjectTable double backed by one real, mutable row. - It honours the batch_processed and status filters, so the poller's compare-and-swap - and the managed-files deletion guard both read the same state a shared Postgres row - would give them. Staleness sweeps (the only queries scoped by created_at) never match. + It honours the batch_processed and status filters and the deletion guard's created_at cutoff, so the + poller's compare-and-swap and the managed-files deletion guard both read the state a shared Postgres row + would give them. The poll lanes' created_at windows are ignored, so every lane sees the row, and each + read hands back a snapshot, so a write that lands mid-cycle is invisible to the row the poller holds. + Writes journal "claim" or "release" when they flip batch_processed, "expire" when they mark the row + stale_expired, and "finalize" otherwise. """ def __init__(self, row: _FakeManagedObjectRow, journal: list): @@ -2956,7 +3120,7 @@ class _FakeManagedObjectTable: def _matches(self, where: dict) -> bool: for key, value in where.items(): if key == "created_at": - if "lt" in value and not self.row.created_at < value["lt"]: + if "gte" in value and not self.row.created_at >= value["gte"]: return False continue if key == "status": @@ -2971,10 +3135,12 @@ class _FakeManagedObjectTable: async def _update_many(self, *, where: dict, data: dict) -> int: if not self._matches(where): return 0 - if "updated_at" in where: + if "batch_processed" in data: + self.journal.append("claim" if data["batch_processed"] else "release") + elif data.get("status") == "stale_expired": + self.journal.append("expire") + else: self.journal.append("finalize") - elif "batch_processed" in where: - self.journal.append("claim" if data.get("batch_processed") else "release") for key, value in data.items(): setattr(self.row, key, value) return 1 @@ -2985,7 +3151,7 @@ class _FakeManagedObjectTable: setattr(self.row, key, value) async def _find_many(self, *, where: dict, take=None, order=None) -> list: - return [self.row] if self._matches(where) else [] + return [copy(self.row)] if self._matches(where) else [] class TestMultiPodBatchCostClaim: @@ -3008,6 +3174,7 @@ class TestMultiPodBatchCostClaim: proxy_logging_obj=proxy_logging_obj, prisma_client=prisma, llm_router=llm_router, + clock=_fixed_clock, ) @staticmethod @@ -3465,7 +3632,7 @@ class TestMultiPodBatchCostClaim: cache = MagicMock() cache.async_get_cache = AsyncMock(return_value=None) cache.async_set_cache = AsyncMock() - guard = PROXY_LiteLLMManagedFiles(internal_usage_cache=cache, prisma_client=prisma) + guard = PROXY_LiteLLMManagedFiles(internal_usage_cache=cache, prisma_client=prisma, clock=_fixed_clock) scheduler = MagicMock() scheduler.get_job.return_value = MagicMock() @@ -3683,12 +3850,15 @@ class TestBatchLineItemsFromThePoller: @classmethod @contextmanager - def _provider(cls, journal: list): - """The provider's batch files by id, and a spend-log writer that journals each write.""" + def _provider(cls, journal: list[str], missing_file_error: Exception | None = None) -> Iterator[None]: + """The provider's batch files by id, and a spend-log writer that journals each write. Reading any other + file raises missing_file_error, or the NotFoundError OpenAI answers with.""" files = {"file-input-1": cls._INPUT, _CLAIM_OUTPUT_FILE_ID: cls._OUTPUT, "file-err-1": cls._ERROR} async def _afile_content(**kwargs): journal.append(f"read {kwargs['file_id']}") + if kwargs["file_id"] not in files and missing_file_error is not None: + raise missing_file_error if kwargs["file_id"] not in files: from litellm.exceptions import NotFoundError @@ -3861,7 +4031,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize", "read file-err-1", "read file-input-1", "lines"] + assert journal == ["claim", "read file-err-1", "read file-input-1", "lines"] (payload,) = logger.payloads assert [(line["custom_id"], line["status"], line["response_cost"]) for line in payload["lines"]] == [ ("ok-1", "failure", 0.0) @@ -3892,7 +4062,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize", "read file-err-1", "read file-input-1", "lines"] + assert journal == ["claim", "read file-err-1", "read file-input-1", "lines"] ((line,),) = [payload["lines"] for payload in logger.payloads] assert line["request"] == json.loads(self._INPUT) @@ -3910,7 +4080,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize", "read file-err-1", "read file-input-1"] + assert journal == ["claim", "read file-err-1", "read file-input-1"] assert logger.payloads == [] @pytest.mark.asyncio @@ -3970,11 +4140,37 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["read file-out-deleted", "finalize", "read file-err-1", "read file-input-1", "lines"] + assert journal == ["read file-out-deleted", "claim", "read file-err-1", "read file-input-1", "lines"] (payload,) = logger.payloads assert [(line["custom_id"], line["status"]) for line in payload["lines"]] == [("ok-1", "failure")] assert row.batch_processed is True + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("status_code", "missing_object", "retired"), + ( + (404, "bucket/job-1/predictions.jsonl", True), + (404, "bucket/job-2/predictions.jsonl", False), + (500, "bucket/job-1/predictions.jsonl", False), + ), + ) + async def test_a_storage_404_for_the_output_object_retires_the_batch_unbilled( + self, status_code: int, missing_object: str, retired: bool + ) -> None: + """Vertex AI reports a deleted gs:// output as a VertexAIError 404 that names the object without its + scheme, and no later poll can read that object either.""" + from litellm.llms.vertex_ai.common_utils import VertexAIError + + journal: list[str] = [] + row, instance = self._setup(journal, output_file_id="gs://bucket/job-1/predictions.jsonl") + missing: Final = VertexAIError(status_code=status_code, message=f"No such object: {missing_object}") + + with self._provider(journal, missing_file_error=missing): + await instance.check_batch_cost() + + assert row.batch_processed is retired + assert "bill" not in journal + @pytest.mark.asyncio async def test_a_terminal_batch_with_no_lines_sends_nothing(self, monkeypatch): import litellm @@ -3988,7 +4184,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize"] + assert journal == ["claim"] assert logger.payloads == [] @pytest.mark.asyncio @@ -4088,16 +4284,19 @@ class TestBatchLineItemsFromThePoller: _, instance = self._setup(journal) monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),)) - async def _an_earlier_batch_waits_to_store_its_lines(): + find_jobs: Final = instance._find_jobs + + async def _an_earlier_batch_waits_to_store_its_lines(where: Mapping[str, object]) -> Sequence[object]: instance._store_batch_line_items_in_background( MagicMock(), MagicMock(), "model-123", "batch-earlier", billed=True ) await asyncio.sleep(0.01) + return await find_jobs(where) async def _store(job, response, model_id, batch_id, *, billed): journal.append(f"store {batch_id}") - monkeypatch.setattr(instance, "_cleanup_stale_managed_objects", _an_earlier_batch_waits_to_store_its_lines) + monkeypatch.setattr(instance, "_find_jobs", _an_earlier_batch_waits_to_store_its_lines) monkeypatch.setattr(instance, "_store_batch_line_items", _store) with self._provider(journal): await instance.check_batch_cost() @@ -4345,23 +4544,261 @@ class TestCreatorTurnedOffMessageLogging: instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called() +_PAST_THE_STALENESS_CUTOFF: Final = timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, microseconds=1) +_INSIDE_THE_STALENESS_CUTOFF: Final = timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) +_PAST_THE_GIVE_UP_AGE: Final = GIVE_UP_POLLING_AFTER + timedelta(microseconds=1) +_INSIDE_THE_GIVE_UP_AGE: Final = GIVE_UP_POLLING_AFTER + + +def _provider_batch(status: str) -> LiteLLMBatch: + return LiteLLMBatch( + id="batch-456", + status=status, + output_file_id=None, + error_file_id=None, + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) + + @pytest.mark.asyncio -async def test_old_completed_batch_remains_available_for_costing() -> None: +async def test_an_unbilled_batch_past_the_staleness_cutoff_no_longer_blocks_deleting_its_input_file() -> None: row: Final = _FakeManagedObjectRow() row.status = "completed" - row.created_at = datetime(2000, 1, 1, tzinfo=timezone.utc) + row.created_at = _NOW - _PAST_THE_STALENESS_CUTOFF prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) - poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) - await poller._cleanup_stale_managed_objects() + await TestMultiPodBatchCostClaim._run_deletion_guard(prisma, "file-input-1") assert row.batch_processed is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("age", "file_field", "file_id"), + ( + (_INSIDE_THE_STALENESS_CUTOFF, "input_file_id", "file-input-1"), + (_PAST_THE_STALENESS_CUTOFF, "output_file_id", _CLAIM_OUTPUT_FILE_ID), + (_PAST_THE_STALENESS_CUTOFF, "error_file_id", "file-error-1"), + ), + ids=("input_inside_the_cutoff", "output_past_the_cutoff", "error_past_the_cutoff"), +) +async def test_an_unbilled_batch_still_blocks_deleting_a_file_its_cost_may_need( + age: timedelta, file_field: str, file_id: str +) -> None: + row: Final = _FakeManagedObjectRow() + row.status = "completed" + row.created_at = _NOW - age + row.file_object = json.dumps({"id": "batch-456", "status": "completed", file_field: file_id}) + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + + with pytest.raises(HTTPException) as refused: + await TestMultiPodBatchCostClaim._run_deletion_guard(prisma, file_id) + + assert refused.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_a_stale_expired_batch_no_longer_blocks_deleting_its_output_file() -> None: + row: Final = _FakeManagedObjectRow() + row.status = "stale_expired" + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, []) + + await TestMultiPodBatchCostClaim._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) + + assert row.status == "stale_expired" + + +@pytest.mark.asyncio +async def test_a_batch_past_the_give_up_age_that_the_provider_finished_is_still_billed() -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - _PAST_THE_GIVE_UP_AGE + journal: Final = [] + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, journal) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, TestMultiPodBatchCostClaim._router()) + + with TestMultiPodBatchCostClaim._billing_patches(journal): + await poller.check_batch_cost() + + assert journal == ["fetch", "claim", "bill"] + assert row.batch_processed is True + + +_A_REJECTED_KEY: Final = litellm.AuthenticationError(message="invalid api key", llm_provider="openai", model="gpt-4") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("unified_object_id", "retrieve"), + ( + (_CLAIM_UNIFIED_BATCH_ID, _provider_batch("in_progress")), + (_CLAIM_UNIFIED_BATCH_ID, RuntimeError("timed out")), + ( + _CLAIM_UNIFIED_BATCH_ID, + litellm.InternalServerError(message="upstream error", llm_provider="openai", model="gpt-4"), + ), + (_CLAIM_UNIFIED_BATCH_ID, litellm.RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4")), + (_CLAIM_UNIFIED_BATCH_ID, _A_REJECTED_KEY), + ("batch-raw-1", _provider_batch("in_progress")), + ), + ids=( + "still_running", + "retrieve_timed_out", + "provider_5xx", + "provider_429", + "key_rejected", + "unroutable", + ), +) +@pytest.mark.parametrize( + ("age", "db_status", "after_the_poll"), + ( + (_PAST_THE_GIVE_UP_AGE, "in_progress", ("stale_expired", False)), + (_PAST_THE_GIVE_UP_AGE, "completed", ("completed", False)), + (_INSIDE_THE_GIVE_UP_AGE, "in_progress", ("in_progress", False)), + (_INSIDE_THE_GIVE_UP_AGE, "completed", ("completed", False)), + ), + ids=("unfinished_past_the_age", "finished_past_the_age", "unfinished_at_the_age", "finished_at_the_age"), +) +async def test_only_an_unfinished_batch_a_poll_cannot_settle_is_given_up_once_past_the_give_up_age( + unified_object_id: str, + retrieve: LiteLLMBatch | Exception, + age: timedelta, + db_status: str, + after_the_poll: tuple[str, bool], +) -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - age + row.status = db_status + row.unified_object_id = unified_object_id + router: Final = TestMultiPodBatchCostClaim._router() + router.aretrieve_batch.side_effect = [retrieve] + poller: Final = TestMultiPodBatchCostClaim._instance(TestMultiPodBatchCostClaim._prisma(row, []), router) + + with TestMultiPodBatchCostClaim._billing_patches([]): + await poller.check_batch_cost() + + assert (row.status, row.batch_processed) == after_the_poll + + +async def _a_vertex_job_that_no_longer_exists( + *, model: str, batch_id: str, litellm_metadata: Mapping[str, object] +) -> LiteLLMBatch: + from litellm.llms.vertex_ai.common_utils import VertexAIError + + raise VertexAIError( + status_code=404, + message='Error: 404 {"error":{"code":404,"message":"The BatchPredictionJob does not exist.","status":"NOT_FOUND"}}', + ) + + +async def _a_finished_batch_whose_output_never_appeared( + *, model: str, batch_id: str, litellm_metadata: Mapping[str, object] +) -> LiteLLMBatch: + return _provider_batch("completed") + + +async def _an_output_file_the_key_cannot_read() -> None: + raise litellm.PermissionDeniedError( + message="forbidden", + llm_provider="openai", + model="gpt-4", + response=httpx.Response(403, request=httpx.Request("GET", "https://api.openai.com/v1/files")), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("db_status", ("in_progress", "completed")) +async def test_a_batch_past_the_give_up_age_whose_output_id_lags_is_billed_once_it_appears(db_status: str) -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - _PAST_THE_GIVE_UP_AGE + row.status = db_status + journal: Final = [] + router: Final = TestMultiPodBatchCostClaim._router() + router.aretrieve_batch.side_effect = [ + _provider_batch("completed"), + _provider_batch("completed"), + router.aretrieve_batch.return_value, + ] + poller: Final = TestMultiPodBatchCostClaim._instance(TestMultiPodBatchCostClaim._prisma(row, journal), router) + + with TestMultiPodBatchCostClaim._billing_patches(journal): + for _ in range(3): + await poller.check_batch_cost() + + assert journal == ["fetch", "claim", "bill"] + assert row.batch_processed is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("retrieve", "during_fetch"), + ( + (_A_REJECTED_KEY, None), + (_a_vertex_job_that_no_longer_exists, None), + (_a_finished_batch_whose_output_never_appeared, None), + (None, _an_output_file_the_key_cannot_read), + ), + ids=("key_rejected", "batch_404_without_its_id", "output_id_pending", "output_unreadable"), +) +async def test_a_finished_batch_a_poll_cannot_bill_past_the_give_up_age_bills_once_the_cause_is_fixed( + retrieve: Exception | Callable[..., Awaitable[LiteLLMBatch]] | None, + during_fetch: Callable[[], Awaitable[None]] | None, +) -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - _PAST_THE_GIVE_UP_AGE + row.status = "completed" + router: Final = TestMultiPodBatchCostClaim._router() + router.aretrieve_batch.side_effect = retrieve + journal: Final = [] + prisma: Final = TestMultiPodBatchCostClaim._prisma(row, journal) + poller: Final = TestMultiPodBatchCostClaim._instance(prisma, router) + + with TestMultiPodBatchCostClaim._billing_patches(journal, during_fetch=during_fetch): + await poller.check_batch_cost() + + assert (row.status, row.batch_processed) == ("completed", False) + with pytest.raises(HTTPException): + await TestMultiPodBatchCostClaim._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) + + router.aretrieve_batch.side_effect = None + with TestMultiPodBatchCostClaim._billing_patches(journal): + await poller.check_batch_cost() + + assert journal[-1] == "bill" + assert row.batch_processed is True + + +@pytest.mark.asyncio +async def test_giving_up_on_a_batch_keeps_a_finished_status_a_client_wrote_meanwhile() -> None: + row: Final = _FakeManagedObjectRow() + row.created_at = _NOW - _PAST_THE_GIVE_UP_AGE + router: Final = TestMultiPodBatchCostClaim._router() + + async def _a_client_reads_the_finished_batch_while_the_poller_retrieve_is_refused( + *, model: str, batch_id: str, litellm_metadata: Mapping[str, object] + ) -> LiteLLMBatch: + row.status = "completed" + raise _A_REJECTED_KEY + + router.aretrieve_batch.side_effect = _a_client_reads_the_finished_batch_while_the_poller_retrieve_is_refused + poller: Final = TestMultiPodBatchCostClaim._instance(TestMultiPodBatchCostClaim._prisma(row, []), router) + + with TestMultiPodBatchCostClaim._billing_patches([]): + await poller.check_batch_cost() + assert row.status == "completed" + assert row.batch_processed is False @pytest.mark.asyncio @pytest.mark.parametrize("progress", ["claimed", "refreshed"]) -async def test_stale_retirement_preserves_concurrent_progress(progress: str) -> None: +async def test_stale_retirement_lands_unless_another_pod_claimed_the_row(progress: str) -> None: + """A client GET that rewrote the row while the poller held it moves updated_at but leaves the batch unbilled, + so the retirement still lands.""" row: Final = _FakeManagedObjectRow() selected: Final = copy(row) if progress == "claimed": @@ -4372,13 +4809,16 @@ async def test_stale_retirement_preserves_concurrent_progress(progress: str) -> await poller._retire_job(selected, "provider batch disappeared") - assert row.batch_processed is (progress == "claimed") - assert row.updated_at > selected.updated_at + assert row.batch_processed is True @pytest.mark.asyncio -@pytest.mark.parametrize("progress", ["claimed", "refreshed"]) -async def test_stale_terminal_response_preserves_concurrent_progress(progress: str) -> None: +@pytest.mark.parametrize( + ("progress", "finalized", "expected_status"), [("claimed", False, "in_progress"), ("refreshed", True, "failed")] +) +async def test_stale_terminal_response_lands_unless_another_pod_claimed_the_row( + progress: str, finalized: bool, expected_status: str +) -> None: from litellm.types.utils import LiteLLMBatch row: Final = _FakeManagedObjectRow() @@ -4399,11 +4839,10 @@ async def test_stale_terminal_response_preserves_concurrent_progress(progress: s status="failed", ) - assert await poller._finalize_unbilled_terminal_job(selected, response) is False + assert await poller._finalize_unbilled_terminal_job(selected, response) is finalized - assert row.file_object == '{"revision":"fresh"}' - assert row.status == "in_progress" - assert row.batch_processed is (progress == "claimed") + assert row.status == expected_status + assert row.batch_processed is True @pytest.mark.asyncio @@ -4447,7 +4886,19 @@ def test_another_resource_in_a_404_is_not_permanent_loss(resource_id: str, other @pytest.mark.parametrize("resource_id", ("batch-123", "file-123", "gs://bucket/job/predictions.jsonl")) -@pytest.mark.parametrize("message", ("No object found with id '{id}'.", 'File "{id}" does not exist', "{id} not found")) +@pytest.mark.parametrize( + "message", + ( + "No object found with id '{id}'.", + 'File "{id}" does not exist', + "{id} not found", + "No such object: {id}.", + "{id}: not found", + "Missing {id}, giving up", + "Not found (id {id})", + "Not found ({id})", + ), +) def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, message: str) -> None: from litellm.exceptions import NotFoundError from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost @@ -4457,6 +4908,23 @@ def test_a_404_naming_the_exact_resource_is_permanent_loss(resource_id: str, mes assert CheckBatchCost._is_output_file_gone_at_provider(error, resource_id) is True +@pytest.mark.parametrize( + ("missing_key", "gone"), (("out/job-123/input.jsonl.out", True), ("out/job-123/other.jsonl.out", False)) +) +def test_an_s3_no_such_key_for_the_output_object_is_permanent_loss(missing_key: str, gone: bool) -> None: + from litellm.llms.bedrock.common_utils import BedrockError + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + # S3 GetObject 404 body shape, checked live against a missing key on 2026-10-20 + error: Final = BedrockError( + status_code=404, + message='\nNoSuchKey' + f"The specified key does not exist.{missing_key}" + "RH", + ) + assert CheckBatchCost._is_output_file_gone_at_provider(error, "s3://qa-bucket/out/job-123/input.jsonl.out") is gone + + @pytest.mark.parametrize("output_file_id", (None, "")) def test_a_404_without_an_output_identity_remains_retryable(output_file_id: str | None) -> None: from litellm.exceptions import NotFoundError diff --git a/tests/unit/proxy/common_utils/test_token_count_settings.py b/tests/unit/proxy/common_utils/test_token_count_settings.py index 366d8b010d0..43edc7e449f 100644 --- a/tests/unit/proxy/common_utils/test_token_count_settings.py +++ b/tests/unit/proxy/common_utils/test_token_count_settings.py @@ -32,22 +32,31 @@ def _in_us_west_2() -> str: ({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)), ( {"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}}, - TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")), + TokenCountSettings( + framing=5, + count_model=TokenCountModel(model="custom-model", region="ap-southeast-4", configured=True), + ), ), ( {"token_count_model": {"model": "custom-model", "litellm_provider": "bedrock_mantle"}}, TokenCountSettings( framing=97, - count_model=TokenCountModel(model="custom-model", region=None, litellm_provider="bedrock_mantle"), + count_model=TokenCountModel( + model="custom-model", region=None, litellm_provider="bedrock_mantle", configured=True + ), ), ), ( {"token_count_model": {"model": "custom-model"}}, - TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + TokenCountSettings( + framing=97, count_model=TokenCountModel(model="custom-model", region=None, configured=True) + ), ), ( {"token_count_model": {"model": "custom-model", "region": " "}}, - TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)), + TokenCountSettings( + framing=97, count_model=TokenCountModel(model="custom-model", region=None, configured=True) + ), ), ({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)), ( @@ -80,26 +89,26 @@ def test_config_overrides_the_cost_map_default_key_by_key( @pytest.mark.parametrize( - ("deployment_region", "configured", "expected_region"), + ("deployment_region", "configured", "expected"), [ - ("us-west-2", {}, "us-east-1"), - ("eu-central-1", {}, None), - ("ca-central-1", {}, None), - ("us-gov-west-1", {}, None), + ("us-west-2", {}, _OPUS_US_EAST_1), + ("eu-central-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), + ("ca-central-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), + ("us-gov-west-1", {}, TokenCountModel(model="anthropic.claude-opus-5-5", region=None)), ( "eu-central-1", {"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}}, - "us-east-1", + TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1", configured=True), ), ], ) def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography( - deployment_region: str, configured: dict[str, object], expected_region: str | None + deployment_region: str, configured: dict[str, object], expected: TokenCountModel ) -> None: settings: Final = resolve_token_count_settings( configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region ) - assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region) + assert settings.count_model == expected @pytest.mark.parametrize("prefix", ("", "bedrock/", "bedrock_mantle/")) @@ -220,3 +229,10 @@ def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected assert hinted.status_code == 404 + + +def test_hint_stays_off_an_error_raised_before_any_provider_call() -> None: + message: Final = "Invalid request for token counting: LiteLLM cannot count a tool of type 'mcp' for this model" + failure: Final = ProviderCountFailure(status_code=400, message=message, rejected_before_provider_call=True) + + assert failure.with_hint("Enable eu-west-1").message == message 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_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/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/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/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index f692144ef78..3e7597e903c 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -383,6 +383,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/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 52816dc50a8..790790459c3 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -10826,12 +10826,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 @@ -10932,6 +10939,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_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 5484bc38317..0e8fc6823bd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -12,6 +12,7 @@ import threading import traceback import uuid from collections.abc import Callable, Coroutine, Mapping +from dataclasses import replace from functools import partial from pathlib import Path from types import MappingProxyType @@ -209,7 +210,6 @@ async def test_gpt_token_counting(): assert response.request_model == "gpt-4" - @pytest.mark.asyncio async def test_token_counter_leaves_the_routers_deployment_without_the_credentials_it_resolves(tmp_path, monkeypatch): """token_counter reads the router's own deployment, so loading the named credential into it would leave the @@ -1685,6 +1685,7 @@ async def _ladder( system=_COUNT_REQUEST.system, ), model_to_use="us.anthropic.claude-sonnet-5-5", + request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=local_estimate_allowed, ) @@ -1748,6 +1749,7 @@ async def test_token_count_model_provider_selects_the_endpoint_and_preserves_the system=_COUNT_REQUEST.system, ), model_to_use="us.anthropic.claude-sonnet-5-5", + request_model=_COUNT_REQUEST.model, settings=settings, local_estimate_allowed=False, ) @@ -1786,15 +1788,22 @@ async def test_token_count_model_counts_the_same_request_in_its_own_region_witho assert _SONNET_DEPLOYMENT["litellm_params"]["aws_region_name"] == "us-west-2" +@pytest.mark.parametrize( + "count_model", + [None, TokenCountModel(model="us.anthropic.claude-sonnet-5-5", region=None)], + ids=["no_token_count_model", "shipped_token_count_model_in_the_deployment_region"], +) @pytest.mark.parametrize("status_code", [401, 403, 429, 500]) @pytest.mark.asyncio -async def test_errors_other_than_cannot_count_raise_even_with_framing(status_code: int): +async def test_errors_other_than_cannot_count_raise_even_with_framing( + count_model: TokenCountModel | None, status_code: int +): counter: Final = _ScriptedCounter( {"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=status_code)} ) with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=10, count_model=None)) + await _ladder(counter, TokenCountSettings(framing=10, count_model=count_model)) assert (raised.value.code, raised.value.message) == (str(status_code), f"scripted {status_code} for sonnet") @@ -1840,28 +1849,97 @@ async def test_uncountable_model_without_framing_raises(): assert raised.value.code == "404" +@pytest.mark.parametrize( + ("count_model", "counted_as"), + [ + (None, "us.anthropic.claude-sonnet-5-5"), + ( + TokenCountModel(model="anthropic.claude-sonnet-5-5", region=None, litellm_provider="bedrock"), + "bedrock/anthropic.claude-sonnet-5-5", + ), + (_OPUS_IN_US_EAST_1, "us.anthropic.claude-opus-5-5"), + ], + ids=[ + "no_token_count_model", + "shipped_token_count_model_in_the_deployment_region", + "shipped_token_count_model_pinned_to_a_region", + ], +) @pytest.mark.asyncio -async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate(): - counter: Final = _ScriptedCounter({"us.anthropic.claude-sonnet-5-5": _provider_reply("sonnet", status_code=400)}) +async def test_uncountable_model_with_framing_falls_through_to_the_local_estimate( + count_model: TokenCountModel | None, counted_as: str +): + counter: Final = _ScriptedCounter({counted_as: _provider_reply("sonnet", status_code=400)}) - result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=None)) + result: Final = await _ladder(counter, TokenCountSettings(framing=0, count_model=count_model)) assert isinstance(result, ProviderCountFailure) - assert len(counter.calls) == 1 + assert [call[0] for call in counter.calls] == [counted_as] -@pytest.mark.parametrize("status_code", [400, 404]) +@pytest.mark.parametrize("status_code", [401, 403]) @pytest.mark.asyncio -async def test_token_count_model_error_stands_even_with_framing(status_code: int): +async def test_framing_stands_in_when_a_shipped_token_count_model_region_denies_the_account(status_code: int): + counter: Final = _ScriptedCounter( + {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} + ) + + result: Final = await _ladder(counter, TokenCountSettings(framing=50, count_model=_OPUS_IN_US_EAST_1)) + + assert isinstance(result, ProviderCountFailure) + assert result.status_code == status_code + + +@pytest.mark.parametrize( + ("count_model", "region_hints"), + [ + (replace(_OPUS_IN_US_EAST_1, configured=True), 1), + (replace(_OPUS_IN_US_EAST_1, region=None, configured=True), 0), + ], + ids=["pinned_to_a_region", "in_the_deployment_region"], +) +@pytest.mark.parametrize("status_code", [400, 401, 403, 404]) +@pytest.mark.asyncio +async def test_configured_token_count_model_error_stands_even_with_framing( + count_model: TokenCountModel, region_hints: int, status_code: int +): counter: Final = _ScriptedCounter( {"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=status_code)} ) with pytest.raises(ProxyException) as raised: - await _ladder(counter, TokenCountSettings(framing=78, count_model=_OPUS_IN_US_EAST_1)) + await _ladder(counter, TokenCountSettings(framing=78, count_model=count_model)) assert raised.value.code == str(status_code) - assert "us.anthropic.claude-opus-5-5 in us-east-1" in raised.value.message + assert raised.value.message.startswith(f"scripted {status_code} for opus") + assert raised.value.message.count("counts tokens with us.anthropic.claude-opus-5-5 in us-east-1") == region_hints + + +@pytest.mark.parametrize( + ("count_model", "region_hints"), + [(_OPUS_IN_US_EAST_1, frozenset({True})), (replace(_OPUS_IN_US_EAST_1, region=None), frozenset())], + ids=["pinned_to_a_region", "in_the_deployment_region"], +) +@pytest.mark.asyncio +async def test_framing_standing_in_for_a_pinned_token_count_model_warns_with_its_model_and_region( + count_model: TokenCountModel, region_hints: frozenset[bool], caplog: pytest.LogCaptureFixture +): + counter: Final = _ScriptedCounter({"us.anthropic.claude-opus-5-5": _provider_reply("opus", status_code=404)}) + verbose_proxy_logger.addHandler(caplog.handler) + try: + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + result: Final = await _ladder(counter, TokenCountSettings(framing=50, count_model=count_model)) + finally: + verbose_proxy_logger.removeHandler(caplog.handler) + + assert isinstance(result, ProviderCountFailure) + stand_in_warnings: Final = [ + r.getMessage() for r in caplog.records if "falls back to the framed local estimate" in r.getMessage() + ] + assert { + f"for {_COUNT_REQUEST.model} answered 404 from us.anthropic.claude-opus-5-5 in us-east-1" in message + for message in stand_in_warnings + } == region_hints @pytest.mark.asyncio @@ -2837,7 +2915,9 @@ async def test_configured_token_count_model_overrides_the_shipped_one_and_an_emp async def test_shipped_token_count_model_keeps_the_prompt_in_the_deployment_geography( monkeypatch: pytest.MonkeyPatch, aws_region_name: str, counted_in: str ): - counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + counter: Final = _ScriptedCounter( + {"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)} + ) monkeypatch.setattr( litellm.proxy.proxy_server, "llm_router", @@ -2888,7 +2968,9 @@ class _VaultWaitingForLoopProgress: async def test_region_lookup_for_a_shipped_count_model_leaves_the_event_loop_free(monkeypatch: pytest.MonkeyPatch): from litellm.proxy import proxy_server - counter: Final = _ScriptedCounter({"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)}) + counter: Final = _ScriptedCounter( + {"bedrock_mantle/anthropic.claude-opus-5-5": _provider_reply("opus", total_tokens=447)} + ) def bedrock_counter(deployment: object, model: str) -> tuple[_ScriptedCounter, str, str]: return counter, model, "bedrock" @@ -3653,7 +3735,11 @@ async def test_responses_count_sends_openai_the_provider_file_id_inference_sends model_list=[ { "model_name": "gpt-files", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://count.test/v1"}, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://count.test/v1", + }, "model_info": {"id": "deployment-a"}, } ] @@ -3833,9 +3919,7 @@ async def test_responses_count_on_claude_rejects_forced_tool_choice_without_drop monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_forced_tool_use": False}) with pytest.raises(ProxyException) as rejected: - await _claude_count( - monkeypatch, {}, {"input": "Weather?", "tools": [_WEATHER_TOOL], "tool_choice": "required"} - ) + await _claude_count(monkeypatch, {}, {"input": "Weather?", "tools": [_WEATHER_TOOL], "tool_choice": "required"}) assert (rejected.value.code, "forced tool use" in rejected.value.message) == ("400", True) 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/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_api_lifecycle.py b/tests/unit/responses/test_responses_api_lifecycle.py index 8805bc97d10..80da62f884a 100644 --- a/tests/unit/responses/test_responses_api_lifecycle.py +++ b/tests/unit/responses/test_responses_api_lifecycle.py @@ -6,6 +6,7 @@ import pytest import respx import litellm +from litellm.types.llms.openai import ResponsesAPIResponse from tests.unit.proxy.conftest import httpx_transport pytestmark: Final = pytest.mark.usefixtures(httpx_transport.__name__) @@ -42,6 +43,12 @@ async def _invoke_cancel_response(sync_mode: bool, call_kwargs: dict[str, object return await litellm.acancel_responses(**call_kwargs) +async def _invoke_get_response(sync_mode: bool, call_kwargs: dict[str, object]) -> object: + if sync_mode: + return litellm.get_responses(**call_kwargs) + return await litellm.aget_responses(**call_kwargs) + + @pytest.mark.parametrize( ("provider", "operation"), ( @@ -102,8 +109,43 @@ async def test_cancel_responses_404_surfaces_openai_api_error_with_exact_url( } with pytest.raises(openai.APIError) as exc_info: await _invoke_cancel_response(sync_mode, call_kwargs) - requests: Final = tuple(route.calls) + requests: Final[tuple[respx.models.Call, ...]] = tuple(route.calls) assert exc_info.value.status_code == 404 assert len(requests) == 1 assert str(requests[0].request.url) == expected_url + + +@pytest.mark.parametrize("sync_mode", (True, False)) +@pytest.mark.asyncio +async def test_cancel_responses_reply_carries_the_id_retrieve_gives_for_the_same_run(sync_mode: bool) -> None: + upstream_id: Final = "resp_upstream123" + reply: Final = { + "id": upstream_id, + "object": "response", + "created_at": 0, + "status": "cancelled", + "model": "gpt-6.1-terra", + "output": [], + } + call_kwargs: Final = { + "custom_llm_provider": "openai", + "api_key": "sk-test", + "response_id": upstream_id, + "litellm_metadata": {"model_info": {"id": "deployment-1"}}, + } + + with respx.mock() as mock_router: + mock_router.post(f"https://api.openai.com/v1/responses/{upstream_id}/cancel").mock( + return_value=httpx.Response(status_code=200, json=reply) + ) + mock_router.get(f"https://api.openai.com/v1/responses/{upstream_id}").mock( + return_value=httpx.Response(status_code=200, json=reply) + ) + cancelled: Final = await _invoke_cancel_response(sync_mode, call_kwargs) + retrieved: Final = await _invoke_get_response(sync_mode, call_kwargs) + + assert isinstance(cancelled, ResponsesAPIResponse) + assert isinstance(retrieved, ResponsesAPIResponse) + assert cancelled.id == retrieved.id + assert cancelled.id != upstream_id 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_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 4299b61c703..0d39016ea79 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -35,6 +35,7 @@ from litellm.router_utils.cooldown_handlers import ( cast_exception_status_to_int, get_cooldown_deployments, mark_advisor_orchestration_failure, + set_cooldown_deployments, should_cooldown_based_on_allowed_fails_policy, ) from litellm.router_utils.fallback_event_handlers import ( @@ -55,6 +56,40 @@ from litellm.types.utils import CooldownDeploymentInfo, CredentialItem from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("exception_type", "expected_cooldown"), + ( + (litellm.APIConnectionError, False), + (litellm.UnsupportedEndpointError, False), + (litellm.BadRequestError, True), + ), +) +async def test_unsupported_endpoint_preserves_cooldown_exclusion_with_bad_request_policy( + exception_type: type[litellm.APIConnectionError | litellm.BadRequestError], expected_cooldown: bool +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "chat", + "litellm_params": {"model": "anthropic/claude-sonnet-5-5", "api_key": "not-needed"}, + "model_info": {"id": "deployment", "allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}, + } + ], + ) + exception: Final = exception_type(message="Request failed", model="claude-sonnet-5-5", llm_provider="anthropic") + set_cooldown_deployments( + litellm_router_instance=router, + original_exception=exception, + exception_status=exception.status_code, + deployment="deployment", + time_to_cooldown=60, + ) + assert await async_get_cooldown_deployments(router, parent_otel_span=None) == ( + ["deployment"] if expected_cooldown else [] + ) + + class TestGetDeploymentCooldownPolicy: def _make_router(self, deployment_id: str, model_info: dict | None = None): router = MagicMock() diff --git a/tests/unit/router_utils/test_get_retry_from_policy.py b/tests/unit/router_utils/test_get_retry_from_policy.py index 42580989a8c..4ff869849dd 100644 --- a/tests/unit/router_utils/test_get_retry_from_policy.py +++ b/tests/unit/router_utils/test_get_retry_from_policy.py @@ -89,6 +89,18 @@ def test_payment_required_error_uses_bad_request_retries(): ) +@pytest.mark.parametrize("default_retries", (None, 0, 3)) +@pytest.mark.parametrize("bad_request_retries", (0, 5)) +def test_unsupported_endpoint_preserves_the_connection_error_retry_budget( + default_retries: int | None, bad_request_retries: int +) -> None: + policy: Final = RetryPolicy(BadRequestErrorRetries=bad_request_retries, DefaultRetries=default_retries) + + assert get_num_retries_from_retry_policy( + exception=_error(litellm.UnsupportedEndpointError), retry_policy=policy + ) == get_num_retries_from_retry_policy(exception=_error(litellm.APIConnectionError), retry_policy=policy) + + @pytest.mark.parametrize("exception_type", (litellm.BadGatewayError,)) def test_default_retries_covers_exceptions_without_a_specific_field(exception_type: type[Exception]): exception: Final = _error(exception_type) diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 3532870df1c..8011ca37cfb 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -15,6 +15,7 @@ from litellm.router_utils.common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + format_content_policy_no_fallback_note, format_fallback_outcome_message, provider_for_generic_call, resolve_model_group_alias, @@ -679,6 +680,14 @@ def test_format_fallback_outcome_preserves_attempted_failure_wording_without_ski ) +def test_content_policy_note_explains_the_failure_when_fallbacks_are_a_plain_model_list() -> None: + note: Final = format_content_policy_no_fallback_note("gpt-bg", ("gpt-bg",), None, ["backup"], False) + + assert "content_policy_fallback=None" in note + assert "LiteLLM: model group 'gpt-bg' failed with the error above. No fallback was attempted." in note + assert "Fallbacks are configured for" not in note + + class TestWarnOnProviderCredentialMismatch: """A deployment that carries one provider's credentials while resolving to another is silently broken: litellm ignores the credentials and sends the 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/secret_managers/test_get_azure_ad_token_provider.py b/tests/unit/secret_managers/test_get_azure_ad_token_provider.py index 4bc7c21d8d2..141072635bc 100644 --- a/tests/unit/secret_managers/test_get_azure_ad_token_provider.py +++ b/tests/unit/secret_managers/test_get_azure_ad_token_provider.py @@ -1,16 +1,25 @@ +import itertools import json import os -from typing import Optional -from unittest.mock import MagicMock, patch +import threading +import time +from collections.abc import Callable, Iterator +from concurrent.futures import Future, ThreadPoolExecutor, wait +from dataclasses import dataclass, field +from queue import SimpleQueue +from typing import Final, Optional +from unittest.mock import ANY, MagicMock, patch # Adds the grandparent directory to sys.path to allow importing project modules import pytest +from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions from azure.core.exceptions import ClientAuthenticationError from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, infer_credential_type_from_environment, + single_flight_bearer_token_provider, ) from litellm.types.secret_managers.get_azure_ad_token_provider import ( AzureCredentialType, @@ -190,7 +199,7 @@ class TestGetAzureAdTokenProvider: tenant_id="test-tenant-id", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -228,7 +237,7 @@ class TestGetAzureAdTokenProvider: client_id="test-client-id" ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -270,7 +279,7 @@ class TestGetAzureAdTokenProvider: certificate_path="/path/to/cert.pem", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -314,7 +323,7 @@ class TestGetAzureAdTokenProvider: password="pwd4cert.pem", ) mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -348,7 +357,7 @@ class TestGetAzureAdTokenProvider: assert callable(result) mock_certificate_credential.assert_called_once_with() mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works @@ -417,9 +426,254 @@ class TestGetAzureAdTokenProvider: assert callable(result) mock_default_azure_credential.assert_called_once_with() mock_get_bearer_token_provider.assert_called_once_with( - mock_credential_instance, "https://cognitiveservices.azure.com/.default" + ANY, "https://cognitiveservices.azure.com/.default" ) # Test that the returned callable works token = result() assert token == "mock-default-token" + + +_CONCURRENT_CALLERS: Final = 4 +_SCOPE: Final = "api://audit/.default" +_YEAR_2100: Final = 4_102_444_800 +_LONG_EXPIRED: Final = 1 +_JOIN_WINDOW_SECONDS: Final = 0.5 +_FAST_FAILING_CALLERS: Final = 16 + + +@dataclass(frozen=True, slots=True) +class _HeldCredential: + expires_on: int = _YEAR_2100 + failing_attempts: int = 0 + attempts: Iterator[int] = field(default_factory=itertools.count) + requests: SimpleQueue[tuple[str, ...]] = field(default_factory=SimpleQueue) + everyone_asked: threading.Event = field(default_factory=threading.Event) + release: threading.Event = field(default_factory=threading.Event) + + def get_token_info(self, *scopes: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + self.requests.put(scopes) + if self.requests.qsize() >= _CONCURRENT_CALLERS: + self.everyone_asked.set() + self.release.wait(5) + attempt: Final = next(self.attempts) + if attempt < self.failing_attempts: + raise ClientAuthenticationError("Entra is down") + return AccessTokenInfo(f"token-{attempt}", self.expires_on) + + +def _call_together(provider: Callable[[], str], credential: _HeldCredential) -> tuple[Future[str], ...]: + start: Final = threading.Barrier(_CONCURRENT_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_CONCURRENT_CALLERS)) + assert credential.requests.get(timeout=5) == (_SCOPE,) + credential.everyone_asked.wait(0.5) + credential.release.set() + return calls + + +def test_concurrent_cold_callers_share_one_token_request() -> None: + credential: Final = _HeldCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + calls: Final = _call_together(provider, credential) + + assert tuple(call.result() for call in calls) == ("token-0",) * _CONCURRENT_CALLERS + assert credential.requests.qsize() == 0 + + +def test_concurrent_callers_after_the_token_expired_share_one_token_request() -> None: + credential: Final = _HeldCredential(expires_on=_LONG_EXPIRED) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + credential.release.set() + provider() + credential.requests.get_nowait() + credential.release.clear() + + calls: Final = _call_together(provider, credential) + + assert tuple(call.result() for call in calls) == ("token-1",) * _CONCURRENT_CALLERS + assert credential.requests.qsize() == 0 + + +def test_a_failed_cold_token_request_fails_only_the_caller_that_made_it() -> None: + credential: Final = _HeldCredential(failing_attempts=1) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + calls: Final = _call_together(provider, credential) + + failed: Final = sorted(isinstance(call.exception(), ClientAuthenticationError) for call in calls) + assert failed == [False] * (_CONCURRENT_CALLERS - 1) + [True] + + +@dataclass(frozen=True, slots=True) +class _ScriptedCredential: + expiries: Iterator[int | None] = field(default_factory=lambda: itertools.repeat(None)) + attempts: Iterator[int] = field(default_factory=itertools.count) + requests: SimpleQueue[None] = field(default_factory=SimpleQueue) + answers: threading.Semaphore = field(default_factory=lambda: threading.Semaphore(0)) + + def get_token_info(self, *_: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + self.requests.put(None) + self.answers.acquire(timeout=5) + attempt: Final = next(self.attempts) + expires_on: Final = next(self.expiries) + if expires_on is None: + raise ClientAuthenticationError("Entra is down") + return AccessTokenInfo(f"token-{attempt}", expires_on) + + +def _answer_together( + provider: Callable[[], str], credential: _ScriptedCredential, token_requests: int +) -> tuple[Future[str], ...]: + start: Final = threading.Barrier(_CONCURRENT_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_CONCURRENT_CALLERS)) + for _ in range(token_requests): + credential.requests.get(timeout=5) + time.sleep(_JOIN_WINDOW_SECONDS) + credential.answers.release() + return calls + + +def test_an_entra_outage_retries_only_the_first_failed_token_request() -> None: + credential: Final = _ScriptedCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + + first_burst: Final = _answer_together(provider, credential, token_requests=2) + later_burst: Final = _answer_together(provider, credential, token_requests=1) + + assert all(isinstance(call.exception(), ClientAuthenticationError) for call in first_burst + later_burst) + assert credential.requests.qsize() == 0 + + +def test_waiting_callers_retry_a_failed_token_request_again_once_entra_recovers() -> None: + credential: Final = _ScriptedCredential(iter((None, None, _LONG_EXPIRED, None, _LONG_EXPIRED))) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + _answer_together(provider, credential, token_requests=2) + credential.answers.release() + assert provider() == "token-2" + credential.requests.get_nowait() + + recovered: Final = _answer_together(provider, credential, token_requests=2) + + assert sorted(call.exception() is None for call in recovered) == [False] + [True] * (_CONCURRENT_CALLERS - 1) + assert {call.result() for call in recovered if call.exception() is None} == {"token-4"} + + +def test_a_failed_refresh_fails_the_request_that_made_it_even_while_a_valid_token_is_held() -> None: + credential: Final = _ScriptedCredential(iter((int(time.time()) + 120, None))) + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + credential.answers.release(2) + + assert provider() == "token-0" + with pytest.raises(ClientAuthenticationError): + provider() + + +def test_callers_waiting_on_a_failed_token_request_share_one_retry_when_entra_fails_fast() -> None: + credential: Final = _ScriptedCredential() + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + start: Final = threading.Barrier(_FAST_FAILING_CALLERS) + + def call() -> str: + start.wait(5) + return provider() + + with ThreadPoolExecutor(max_workers=_FAST_FAILING_CALLERS) as pool: + calls: Final = tuple(pool.submit(call) for _ in range(_FAST_FAILING_CALLERS)) + credential.requests.get(timeout=5) + time.sleep(_JOIN_WINDOW_SECONDS) + credential.answers.release(_FAST_FAILING_CALLERS) + + assert all(isinstance(call.exception(), ClientAuthenticationError) for call in calls) + assert credential.requests.qsize() == 1 + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokens: + first_lifetime_seconds: int = 120 + requests: SimpleQueue[None] = field(default_factory=SimpleQueue) + refresh_started: threading.Event = field(default_factory=threading.Event) + release: threading.Event = field(default_factory=threading.Event) + + def next_token(self) -> tuple[str, int]: + self.requests.put(None) + if self.requests.qsize() == 1: + return "expiring-token", int(time.time()) + self.first_lifetime_seconds + self.refresh_started.set() + self.release.wait(5) + return "refreshed-token", _YEAR_2100 + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokenInfoCredential: + tokens: _ExpiringSoonTokens = field(default_factory=_ExpiringSoonTokens) + + def get_token_info(self, *_: str, options: TokenRequestOptions | None = None) -> AccessTokenInfo: + return AccessTokenInfo(*self.tokens.next_token()) + + +@dataclass(frozen=True, slots=True) +class _ExpiringSoonTokenOnlyCredential: + tokens: _ExpiringSoonTokens = field(default_factory=_ExpiringSoonTokens) + + def get_token( + self, + *_: str, + claims: str | None = None, + tenant_id: str | None = None, + enable_cae: bool = False, + **kwargs: object, + ) -> AccessToken: + return AccessToken(*self.tokens.next_token()) + + +@pytest.mark.parametrize("credential_type", (_ExpiringSoonTokenInfoCredential, _ExpiringSoonTokenOnlyCredential)) +def test_callers_during_a_refresh_keep_the_valid_token_while_one_of_them_refreshes( + credential_type: type[_ExpiringSoonTokenInfoCredential | _ExpiringSoonTokenOnlyCredential], +) -> None: + credential: Final = credential_type() + tokens: Final = credential.tokens + provider: Final = single_flight_bearer_token_provider(credential, _SCOPE) + assert provider() == "expiring-token" + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + refresh: Final = pool.submit(provider) + assert tokens.refresh_started.wait(5) + others: Final = tuple(pool.submit(provider) for _ in range(_CONCURRENT_CALLERS - 1)) + try: + assert tuple(other.result(timeout=2) for other in others) == ("expiring-token",) * (_CONCURRENT_CALLERS - 1) + finally: + tokens.release.set() + + assert (refresh.result(), provider()) == ("refreshed-token", "refreshed-token") + assert tokens.requests.qsize() == 2 + + +def test_callers_during_a_refresh_wait_for_the_new_token_when_the_valid_one_is_about_to_expire() -> None: + tokens: Final = _ExpiringSoonTokens(first_lifetime_seconds=30) + provider: Final = single_flight_bearer_token_provider(_ExpiringSoonTokenInfoCredential(tokens), _SCOPE) + assert provider() == "expiring-token" + + with ThreadPoolExecutor(max_workers=_CONCURRENT_CALLERS) as pool: + refresh: Final = pool.submit(provider) + assert tokens.refresh_started.wait(5) + others: Final = tuple(pool.submit(provider) for _ in range(_CONCURRENT_CALLERS - 1)) + try: + assert not wait(others, timeout=_JOIN_WINDOW_SECONDS).done + finally: + tokens.release.set() + + assert tuple(call.result() for call in (refresh, *others)) == ("refreshed-token",) * _CONCURRENT_CALLERS + assert tokens.requests.qsize() == 2 diff --git a/tests/unit/secret_managers/test_secret_managers_main.py b/tests/unit/secret_managers/test_secret_managers_main.py index 7c6e1de6949..862658d9169 100644 --- a/tests/unit/secret_managers/test_secret_managers_main.py +++ b/tests/unit/secret_managers/test_secret_managers_main.py @@ -2,7 +2,11 @@ import base64 import json import logging import os +import threading import time +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from queue import SimpleQueue from typing import Final from unittest.mock import Mock, patch @@ -12,6 +16,7 @@ import litellm from litellm._logging import redact_internal_details_from_client_message from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.secret_managers.main import ( + _cached_azure_oidc_token_provider, get_secret, normalize_nonempty_secret_str, resolve_local_oidc_reference, @@ -313,6 +318,41 @@ def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider, monkeypat mock_token_provider.assert_called_once_with() +def test_concurrent_azure_oidc_reads_build_one_token_provider(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("AZURE_FEDERATED_TOKEN_FILE", raising=False) + callers: Final = 8 + start: Final = threading.Barrier(callers) + builds: Final = SimpleQueue[str | None]() + everyone_building: Final = threading.Event() + release: Final = threading.Event() + + def held_build(azure_scope: str | None = None) -> Callable[[], str]: + builds.put(azure_scope) + if builds.qsize() >= callers: + everyone_building.set() + release.wait(5) + return lambda: "identity-token" + + def read() -> str | None: + start.wait(5) + return get_secret("oidc/azure/api://concurrent-audience") + + _cached_azure_oidc_token_provider.cache_clear() + with ( + patch("litellm.secret_managers.main.get_azure_ad_token_provider", side_effect=held_build), + ThreadPoolExecutor(max_workers=callers) as pool, + ): + calls: Final = tuple(pool.submit(read) for _ in range(callers)) + assert builds.get(timeout=5) == "api://concurrent-audience" + everyone_building.wait(0.5) + release.set() + tokens: Final = tuple(call.result(timeout=5) for call in calls) + _cached_azure_oidc_token_provider.cache_clear() + + assert tokens == ("identity-token",) * callers + assert builds.qsize() == 0 + + def test_oidc_file_success(tmp_path, monkeypatch): token_file = tmp_path / "token.txt" token_file.write_text("file_token") 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_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 f12129af1fc..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", @@ -9780,6 +10104,37 @@ async def test_atranscription_answers_a_non_string_provider_hint_with_a_bad_requ assert not respx_mock.calls +def test_transcription_rejects_an_unsupported_provider_without_an_http_call(respx_mock: respx.MockRouter) -> None: + model: Final = "test-transcription-model" + provider: Final = "anthropic" + + with pytest.raises(litellm.UnsupportedEndpointError) as raised: + litellm.transcription(model=f"{provider}/{model}", file=("speech.wav", b"RIFF-audio", "audio/wav")) + + assert raised.value.status_code == 400 + assert raised.value.message == f"litellm.BadRequestError: {provider}/{model} does not support audio transcription." + assert raised.value.model == model + assert raised.value.llm_provider == provider + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_atranscription_rejects_an_unsupported_provider_without_an_http_call( + respx_mock: respx.MockRouter, +) -> None: + model: Final = "test-transcription-model" + provider: Final = "anthropic" + + with pytest.raises(litellm.UnsupportedEndpointError) as raised: + await litellm.atranscription(model=f"{provider}/{model}", file=("speech.wav", b"RIFF-audio", "audio/wav")) + + assert raised.value.status_code == 400 + assert raised.value.message == f"litellm.BadRequestError: {provider}/{model} does not support audio transcription." + assert raised.value.model == model + assert raised.value.llm_provider == provider + assert not respx_mock.calls + + @pytest.mark.asyncio async def test_aspeech_infers_the_provider_from_base_url(respx_mock: respx.MockRouter) -> None: base_url: Final = "https://api.groq.com/openai/v1" 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_router/test_router.py b/tests/unit/test_router/test_router.py index c6d02dddae0..b4a3d207bd4 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -3774,6 +3774,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 @@ -7037,6 +7163,307 @@ async def test_aresponses_fallback_attempt_does_not_replay_held_lifecycle_events assert outcome == [fallback_created, fallback_error] +_BACKGROUND_PRIMARY_URL: Final = "https://api.openai.com/v1/responses" +_BACKGROUND_BACKUP_URL: Final = "https://backup.openai.test/v1/responses" + + +def _background_sse_frame(event: Mapping[str, object]) -> bytes: + return f"data: {json.dumps(event)}\n\n".encode() + + +def _background_response_body(status: str, background: bool) -> dict[str, object]: + return { + "id": "resp_upstream_run", + "object": "response", + "created_at": 0, + "status": status, + "model": "gpt-6.1-terra", + "background": background, + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} if status == "completed" else None, + } + + +def _paused_background_stream( + output_may_start: asyncio.Event, + upstream_paused: asyncio.Event, + background: bool, + after_pause: tuple[dict[str, object], ...], +) -> httpx.Response: + async def body() -> AsyncIterator[bytes]: + yield _background_sse_frame( + { + "type": "response.created", + "sequence_number": 0, + "response": _background_response_body("queued", background), + } + ) + yield _background_sse_frame( + { + "type": "response.queued", + "sequence_number": 1, + "response": _background_response_body("queued", background), + } + ) + upstream_paused.set() + await output_may_start.wait() + for event in after_pause: + yield _background_sse_frame(event) + + return httpx.Response(200, content=body(), headers={"content-type": "text/event-stream"}) + + +def _background_output_then_completed(background: bool) -> tuple[dict[str, object], ...]: + return ( + {"type": "response.output_text.delta", "sequence_number": 2, "item_id": "msg_1", "delta": "41"}, + { + "type": "response.completed", + "sequence_number": 3, + "response": _background_response_body("completed", background), + }, + ) + + +def _background_router( + model_name: str, + litellm_params: dict[str, object], + fallbacks: list[dict[str, list[str]]] | None, + content_policy_fallbacks: list[dict[str, list[str]]] | None = None, +) -> Router: + return litellm.Router( + model_list=[ + {"model_name": model_name, "litellm_params": litellm_params}, + { + "model_name": "backup", + "litellm_params": { + "model": "openai/gpt-6.1-terra", + "api_key": "backup-key", + "api_base": "https://backup.openai.test/v1", + }, + }, + ], + fallbacks=fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + num_retries=0, + ) + + +async def _reached_client_before_upstream_paused( + first: "asyncio.Future[object]", upstream_paused: asyncio.Event +) -> bool: + paused: Final = asyncio.ensure_future(upstream_paused.wait()) + done, _ = await asyncio.wait({first, paused}, return_when=asyncio.FIRST_COMPLETED) + paused.cancel() + return first in done + + +@pytest.mark.parametrize( + ("fallbacks", "request_kwargs"), + [ + pytest.param(None, {}, id="no-fallback-configured"), + pytest.param([{"gpt-bg": ["backup"]}], {"disable_fallbacks": True}, id="fallbacks-disabled-for-the-request"), + ], +) +@pytest.mark.asyncio +async def test_aresponses_background_stream_hands_the_response_id_to_the_client_before_output( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + fallbacks: list[dict[str, list[str]]] | None, + request_kwargs: dict[str, bool], +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + output_may_start: Final = asyncio.Event() + upstream_paused: Final = asyncio.Event() + respx_mock.post(_BACKGROUND_PRIMARY_URL).mock( + return_value=_paused_background_stream( + output_may_start, upstream_paused, True, _background_output_then_completed(True) + ) + ) + router: Final = _background_router("gpt-bg", {"model": "openai/gpt-6.1-terra", "api_key": "k"}, fallbacks) + + stream: Final = await router.aresponses(model="gpt-bg", input="hi", background=True, stream=True, **request_kwargs) + first: Final = asyncio.ensure_future(anext(stream)) + arrived_before_output: Final = await _reached_client_before_upstream_paused(first, upstream_paused) + output_may_start.set() + created: Final = await first + rest: Final = [event async for event in stream] + + assert arrived_before_output is True + assert created.type == "response.created" + assert [event.type for event in rest] == [ + "response.queued", + "response.output_text.delta", + "response.completed", + ] + assert rest[-1].response.id == created.response.id + + +@pytest.mark.parametrize( + ("litellm_params", "background", "fallbacks"), + [ + pytest.param({"model": "openai/gpt-6.1-terra", "api_key": "k"}, False, None, id="foreground-request"), + pytest.param( + {"model": "openai/gpt-6.1-terra", "api_key": "k"}, True, [{"gpt-bg": ["backup"]}], id="fallback-configured" + ), + pytest.param( + { + "model": "azure/gpt-6.1-terra", + "api_key": "k", + "api_base": "https://bg-resource.openai.azure.com", + "api_version": "preview", + }, + True, + None, + id="azure-deployment", + ), + ], +) +@pytest.mark.asyncio +async def test_aresponses_stream_keeps_holding_lifecycle_events_until_output_off_the_eligible_background_route( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + litellm_params: dict[str, object], + background: bool, + fallbacks: list[dict[str, list[str]]] | None, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + output_may_start: Final = asyncio.Event() + upstream_paused: Final = asyncio.Event() + respx_mock.route(method="POST", path__regex=r"/responses$").mock( + return_value=_paused_background_stream( + output_may_start, upstream_paused, background, _background_output_then_completed(background) + ) + ) + router: Final = _background_router("gpt-bg", litellm_params, fallbacks) + + stream: Final = await router.aresponses(model="gpt-bg", input="hi", background=background, stream=True) + first: Final = asyncio.ensure_future(anext(stream)) + arrived_before_output: Final = await _reached_client_before_upstream_paused(first, upstream_paused) + output_may_start.set() + events: Final = [await first, *[event async for event in stream]] + + assert arrived_before_output is False + assert upstream_paused.is_set() + assert [event.type for event in events] == [ + "response.created", + "response.queued", + "response.output_text.delta", + "response.completed", + ] + + +@pytest.mark.asyncio +async def test_aresponses_background_stream_surfaces_a_later_error_instead_of_a_fallback_once_its_id_is_out( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + output_may_start: Final = asyncio.Event() + upstream_paused: Final = asyncio.Event() + backup_output_may_start: Final = asyncio.Event() + backup_output_may_start.set() + server_error: Final = { + "type": "error", + "sequence_number": 2, + "error": {"type": "server_error", "code": "server_error", "message": "upstream died", "param": None}, + } + primary_route: Final = respx_mock.post(_BACKGROUND_PRIMARY_URL).mock( + return_value=_paused_background_stream(output_may_start, upstream_paused, True, (server_error,)) + ) + backup_route: Final = respx_mock.post(_BACKGROUND_BACKUP_URL).mock( + return_value=_paused_background_stream( + backup_output_may_start, asyncio.Event(), True, _background_output_then_completed(True) + ) + ) + router: Final = _background_router("gpt-bg", {"model": "openai/gpt-6.1-terra", "api_key": "k"}, None) + + stream: Final = await router.aresponses(model="gpt-bg", input="hi", background=True, stream=True) + first: Final = asyncio.ensure_future(anext(stream)) + arrived_before_output: Final = await _reached_client_before_upstream_paused(first, upstream_paused) + router.fallbacks = [{"gpt-bg": ["backup"]}] + output_may_start.set() + created: Final = await first + rest: Final = [item async for item in _events_until_error(stream)] + + assert arrived_before_output is True + assert created.type == "response.created" + assert [item.type for item in rest[:-1]] == ["response.queued"] + assert isinstance(rest[-1], litellm.InternalServerError) + assert "upstream died" in str(rest[-1]) + assert primary_route.call_count == 1 + assert backup_route.call_count == 0 + + + +_BACKGROUND_TEXT_DELTA: Final = { + "type": "response.output_text.delta", + "sequence_number": 2, + "item_id": "msg_1", + "delta": "41", +} + + +@pytest.mark.parametrize( + ("fallbacks", "content_policy_fallbacks", "output_before_refusal", "router_note_expected"), + [ + pytest.param(None, None, (), True, id="no-fallback-configured"), + pytest.param([{"other-group": ["backup"]}], None, (), True, id="fallbacks-for-another-group"), + pytest.param(None, [{"other-group": ["backup"]}], (), True, id="content-policy-fallbacks-for-another-group"), + pytest.param(None, None, (_BACKGROUND_TEXT_DELTA,), False, id="refused-after-output-text"), + ], +) +@pytest.mark.asyncio +async def test_aresponses_background_stream_content_policy_error_reads_the_same_once_its_id_is_out( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + fallbacks: list[dict[str, list[str]]] | None, + content_policy_fallbacks: list[dict[str, list[str]]] | None, + output_before_refusal: tuple[dict[str, object], ...], + router_note_expected: bool, +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + released_output_may_start: Final = asyncio.Event() + released_upstream_paused: Final = asyncio.Event() + held_output_may_start: Final = asyncio.Event() + held_output_may_start.set() + refusal: Final = { + "type": "error", + "sequence_number": 2 + len(output_before_refusal), + "error": { + "type": "invalid_request_error", + "code": "content_policy_violation", + "message": "Your request was rejected as a result of our safety system.", + "param": None, + }, + } + after_pause: Final = (*output_before_refusal, refusal) + primary_route: Final = respx_mock.post(_BACKGROUND_PRIMARY_URL).mock( + side_effect=[ + _paused_background_stream(released_output_may_start, released_upstream_paused, True, after_pause), + _paused_background_stream(held_output_may_start, asyncio.Event(), False, after_pause), + ] + ) + backup_route: Final = respx_mock.post(_BACKGROUND_BACKUP_URL) + router: Final = _background_router( + "gpt-bg", {"model": "openai/gpt-6.1-terra", "api_key": "k"}, fallbacks, content_policy_fallbacks + ) + + released: Final = await router.aresponses(model="gpt-bg", input="hi", background=True, stream=True) + first: Final = asyncio.ensure_future(anext(released)) + arrived_before_output: Final = await _reached_client_before_upstream_paused(first, released_upstream_paused) + released_output_may_start.set() + await first + released_error: Final = [item async for item in _events_until_error(released)][-1] + held: Final = await router.aresponses(model="gpt-bg", input="hi", stream=True) + held_error: Final = [item async for item in _events_until_error(held)][-1] + + assert arrived_before_output is True + assert isinstance(released_error, litellm.ContentPolicyViolationError) + assert isinstance(held_error, litellm.ContentPolicyViolationError) + assert ("LiteLLM: model group 'gpt-bg' failed with the error above" in held_error.message) is router_note_expected + assert released_error.message == held_error.message + assert primary_route.call_count == 2 + assert backup_route.call_count == 0 + @pytest.mark.asyncio async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" @@ -17296,9 +17723,9 @@ def _anthropic_messages_two_order_primary_model_list() -> list: pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"), ], ) -def test_anthropic_messages_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): +def test_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): router = _anthropic_messages_make_router(**router_kwargs) - assert router._anthropic_messages_stream_can_fall_back("primary", request_kwargs) is expected + assert router._stream_can_fall_back("primary", request_kwargs) is expected @pytest.mark.parametrize( @@ -17308,7 +17735,7 @@ def test_anthropic_messages_stream_can_fall_back_direct_call(router_kwargs, requ pytest.param([1, 1], False, id="same-order-cannot-fall-back"), ], ) -def test_anthropic_messages_stream_can_fall_back_order_levels(orders, expected): +def test_stream_can_fall_back_order_levels(orders, expected): router = Router( model_list=[ { @@ -17319,7 +17746,7 @@ def test_anthropic_messages_stream_can_fall_back_order_levels(orders, expected): ], fallbacks=None, ) - assert router._anthropic_messages_stream_can_fall_back("primary", {"model": "primary"}) is expected + assert router._stream_can_fall_back("primary", {"model": "primary"}) is expected @pytest.mark.parametrize( @@ -17335,12 +17762,12 @@ def test_anthropic_messages_stream_can_fall_back_order_levels(orders, expected): ), ], ) -def test_anthropic_messages_stream_can_fall_back_order_target(request_kwargs, expected): +def test_stream_can_fall_back_order_target(request_kwargs, expected): router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None) - assert router._anthropic_messages_stream_can_fall_back("primary", request_kwargs) is expected + assert router._stream_can_fall_back("primary", request_kwargs) is expected -def test_anthropic_messages_order_levels_direct_call(): +def test_stream_order_levels_direct_call(): router = Router( model_list=[ { @@ -17351,7 +17778,7 @@ def test_anthropic_messages_order_levels_direct_call(): ], fallbacks=None, ) - assert router._anthropic_messages_order_levels("primary", {"model": "primary"}) == (1, 2) + assert router._stream_order_levels("primary", {"model": "primary"}) == (1, 2) @pytest.mark.asyncio @@ -20329,6 +20756,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 @@ -21616,6 +22074,116 @@ def test_router_retry_policy_400_moves_off_the_refusing_deployment_on_every_sync assert accepting.call_count == 1 +_TRANSCRIPTION_AUDIO: Final = ("speech.wav", b"RIFF", "audio/wav") +_TRANSCRIPTION_API_BASE: Final = "https://transcriber.local/v1" + + +def _transcription_deployment( + model_name: str, deployment_id: str, litellm_model: str, weight: int +) -> DeploymentTypedDict: + return { + "model_name": model_name, + "litellm_params": { + "model": litellm_model, + "api_key": "sk-fake", + "api_base": _TRANSCRIPTION_API_BASE, + "weight": weight, + }, + "model_info": {"id": deployment_id}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("policy_scope", ["none", "router", "model_group", "request"]) +@pytest.mark.parametrize( + "caller_credentials", + ({}, {"api_key": "caller-key"}, {"api_base": _TRANSCRIPTION_API_BASE}, {"base_url": _TRANSCRIPTION_API_BASE}), + ids=("configured", "api_key", "api_base", "base_url"), +) +async def test_router_retries_a_transcription_off_a_deployment_whose_provider_cannot_transcribe( + monkeypatch: pytest.MonkeyPatch, + policy_scope: Literal["none", "router", "model_group", "request"], + caller_credentials: Mapping[str, str], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + policy: Final = RetryPolicy(BadRequestErrorRetries=0) + router: Final = litellm.Router( + model_list=[ + _transcription_deployment("transcriber", "cannot-transcribe", "anthropic/claude-sonnet-5-5", weight=1), + _transcription_deployment("transcriber", "can-transcribe", "openai/gpt-transcribe", weight=0), + ], + num_retries=1, + retry_policy=policy if policy_scope == "router" else None, + model_group_retry_policy={"transcriber": policy} if policy_scope == "model_group" else None, + disable_cooldowns=True, + ) + request_kwargs: Final = {"model_group_retry_policy": {"transcriber": policy}} if policy_scope == "request" else {} + + with respx.mock as respx_mock: + upstream: Final = respx_mock.post(f"{_TRANSCRIPTION_API_BASE}/audio/transcriptions").mock( + return_value=httpx.Response(200, json={"text": "hi"}) + ) + response: Final = await router.atranscription( + model="transcriber", file=_TRANSCRIPTION_AUDIO, **request_kwargs, **caller_credentials + ) + + assert response.text == "hi" + assert upstream.call_count == 1 + + +@pytest.mark.parametrize("deployment_count", [1, 2, 3]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "caller_credentials", + ({}, {"api_key": "caller-key"}, {"api_base": _TRANSCRIPTION_API_BASE}, {"base_url": _TRANSCRIPTION_API_BASE}), + ids=("configured", "api_key", "api_base", "base_url"), +) +async def test_router_tries_each_deployment_that_cannot_transcribe_once_then_answers_400( + deployment_count: int, caller_credentials: Mapping[str, str] +) -> None: + router: Final = litellm.Router( + model_list=[ + _transcription_deployment("transcriber", f"cannot-transcribe-{index}", "anthropic/claude-sonnet-5-5", 1) + for index in range(deployment_count) + ], + num_retries=5, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock, pytest.raises(litellm.UnsupportedEndpointError) as raised: + await router.atranscription(model="transcriber", file=_TRANSCRIPTION_AUDIO, **caller_credentials) + + assert raised.value.status_code == 400 + assert router.fail_calls["transcriber"] == deployment_count + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_router_falls_back_from_a_group_that_cannot_transcribe_when_unsupported_params_do_not( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + _transcription_deployment("chat-only", "cannot-transcribe", "anthropic/claude-sonnet-5-5", weight=1), + _transcription_deployment("transcriber", "can-transcribe", "openai/gpt-transcribe", weight=1), + ], + fallbacks=[{"chat-only": ["transcriber"]}], + fallback_on_unsupported_params=False, + num_retries=0, + disable_cooldowns=True, + ) + + with respx.mock as respx_mock: + upstream: Final = respx_mock.post(f"{_TRANSCRIPTION_API_BASE}/audio/transcriptions").mock( + return_value=httpx.Response(200, json={"text": "hi"}) + ) + response: Final = await router.atranscription(model="chat-only", file=_TRANSCRIPTION_AUDIO) + + assert response.text == "hi" + assert upstream.call_count == 1 + + def _make_failure_logging_obj(): return LiteLLMLogging( model="gpt-5.6", @@ -27633,38 +28201,89 @@ def test_router_get_model_group_info() -> None: assert model_group_info.output_cost_per_token > 0 -def test_router_rpm_pre_call_check(): - """ - - for a given model not in model cost map - - with rpm set - - check if rpm check is run - """ - try: - model_list = [ +def test_pre_call_checks_keep_deployment_whose_rpm_is_reached(): + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 0}, + "model_info": {"id": "rpm-zero"}, + }, + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"}, + "model_info": {"id": "uncapped"}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + kept: Final = router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "Hey, how's it going?"}], + ) + + assert [deployment["model_info"]["id"] for deployment in kept] == ["rpm-zero", "uncapped"] + + +@pytest.mark.asyncio +async def test_enable_pre_call_checks_alone_serves_past_deployment_rpm(): + router: Final = Router( + model_list=[ { - "model_name": "fake-openai-endpoint", # openai model name - "litellm_params": { # params for litellm completion/embedding call - "model": "openai/my-fake-model", - "api_key": "my-fake-key", - "api_base": "https://openai-function-calling-workers.tasslexyz.workers.dev/", - "rpm": 0, - }, - }, - ] + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 1}, + "model_info": {"id": "rpm-one"}, + } + ], + enable_pre_call_checks=True, + num_retries=0, + ) - router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0) # type: ignore + served: Final = [ + await router.acompletion( + model="fake-openai-endpoint", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_response="served", + ) + for _ in range(3) + ] - try: - router._pre_call_checks( + assert [response.choices[0].message.content for response in served] == ["served"] * 3 + + +@pytest.mark.asyncio +async def test_enforce_model_rate_limits_returns_429_at_deployment_rpm(): + frozen_minute: Final = datetime(2026, 1, 1, 12, 30, 15, tzinfo=timezone.utc) + router: Final = Router( + model_list=[ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 1}, + "model_info": {"id": "rpm-one"}, + } + ], + enable_pre_call_checks=True, + optional_pre_call_checks=["enforce_model_rate_limits"], + num_retries=0, + ) + + with patch( + "litellm.router_utils.pre_call_checks.model_rate_limit_check.get_utc_datetime", return_value=frozen_minute + ): + first: Final = await router.acompletion( + model="fake-openai-endpoint", + messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_response="served", + ) + with pytest.raises(litellm.RateLimitError) as rate_limited: + await router.acompletion( model="fake-openai-endpoint", - healthy_deployments=model_list, messages=[{"role": "user", "content": "Hey, how's it going?"}], + mock_response="served", ) - pytest.fail("Expected this to fail") - except Exception: - pass - except Exception as e: - pytest.fail(f"Got unexpected exception on router! - {str(e)}") + + assert first.choices[0].message.content == "served" + assert rate_limited.value.status_code == 429 @pytest.mark.parametrize( @@ -29494,26 +30113,6 @@ def test_pre_call_checks_filters_deployment_when_token_count_exceeds_limit(): assert [d["model_info"]["id"] for d in filtered] == ["large"] -def test_pre_call_checks_rpm_check_still_enforced_when_token_counting_fails(): - from litellm.types.router import RouterRateLimitErrorBasic - - model_list: Final = [ - { - "model_name": "fake-openai-endpoint", - "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 0}, - "model_info": {"max_input_tokens": 100}, - }, - ] - router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) - - with pytest.raises(RouterRateLimitErrorBasic): - router._pre_call_checks( - model="fake-openai-endpoint", - healthy_deployments=model_list, - messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, - ) - - def test_pre_call_checks_context_window_check_is_skipped_but_not_fatal_when_token_counting_fails(): model_list: Final = [ { 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..9ac2c923632 100644 --- a/tests/unit/test_spend_logs_guard.py +++ b/tests/unit/test_spend_logs_guard.py @@ -3,6 +3,8 @@ 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 @@ -35,6 +37,13 @@ class TestSpendLogHits: elsewhere: Final = _diff("litellm/proxy/utils.py", 'table = "LiteLLM_SpendLogs"', 5) assert guard.spend_log_hits(removed + other + elsewhere) == () + def test_a_table_whose_name_only_starts_with_the_watched_one_is_not_a_hit(self): + archive: Final = _diff("schema.prisma", "model LiteLLM_SpendLogsArchive {", 900) + index: Final = _diff("schema.prisma", " @@index([team_id], map: \"LiteLLM_SpendLogs_team_id_idx\")", 760) + assert guard.spend_log_hits(archive + index) == ( + '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 @@ -42,6 +51,10 @@ class TestSpendLogHits: assert guard.spend_log_hits(diff) == ("litellm/proxy/schema.prisma: a LiteLLM_SpendLogs",) +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: def test_the_latest_review_of_the_approver_decides(self): reviews: Final = tuple( @@ -55,12 +68,62 @@ 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)) @@ -75,5 +138,15 @@ class TestExitCode: 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"}]') + 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): + monkeypatch.setattr(guard, "added_diff", lambda base, head: _diff("schema.prisma", " x LiteLLM_SpendLogs", 7)) + monkeypatch.setattr(guard, "reviewed_hits", lambda base, commit: () if commit == "old" else None, raising=False) + 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/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/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)/mcp-servers/_components/mcp_connect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx index 81a3994070e..d609c9f2527 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_connect.tsx @@ -5,6 +5,7 @@ import React, { useId, useState } from "react"; import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; import { Button } from "@/components/ui/button"; import { Card, CardContent } from "@/components/ui/card"; +import { Kbd } from "@/components/ui/kbd"; import { Label } from "@/components/ui/label"; import { Switch } from "@/components/ui/switch"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -369,8 +370,7 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = []
- Use the keyboard shortcut ⇧+⌘+J (Mac) or{" "} - Ctrl+Shift+J (Windows/Linux) + Use the keyboard shortcut ⇧+⌘+J (Mac) or Ctrl+Shift+J (Windows/Linux) @@ -380,9 +380,8 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] - Copy the JSON configuration below and paste it into Cursor, then save with{" "} - Cmd+S or{" "} - Ctrl+S + Copy the JSON configuration below and paste it into Cursor, then save with Cmd+S or{" "} + Ctrl+S } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx index db576ed03a5..ec4ac9a8b88 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx @@ -12,6 +12,7 @@ import { getGuardrailsList, testPoliciesAndGuardrails } from "@/components/netwo import PolicySelector, { getPolicyOptionEntries } from "@/components/policies/PolicySelector"; import { Policy } from "@/components/policies/types"; import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion"; +import { Kbd } from "@/components/ui/kbd"; import { AlertTriangle, BarChart3, @@ -1306,10 +1307,7 @@ export default function ComplianceUI({ />
- Press Enter to - submit ·{" "} - Shift+Enter for new - line + Press Enter to submit · Shift+Enter for new line {quickTestInput.length}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx index 50ffb89ae83..c7eb8aaec48 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/policies/_components/ai_suggestion_modal.tsx @@ -5,6 +5,7 @@ import { Card, CardContent } from "@/components/ui/card"; import { Checkbox } from "@/components/ui/checkbox"; import { Dialog, DialogContent, DialogTitle } from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; +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"; @@ -654,8 +655,7 @@ const AiSuggestionModal: React.FC = ({ />
- Press Enter to - submit + Press Enter to submit
diff --git a/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx b/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx index 9b23dbebd6c..38bb987b7c4 100644 --- a/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx +++ b/ui/litellm-dashboard/src/components/CommandPalette/CommandPalette.tsx @@ -7,6 +7,7 @@ import { usePathname, useRouter } from "next/navigation"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import { Dialog, DialogContent, DialogTitle } from "@/components/ui/dialog"; +import { Kbd } from "@/components/ui/kbd"; import { useTerminology } from "@/contexts/TerminologyContext"; import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants"; import { keyDetailHref } from "@/utils/entityLinks"; @@ -31,9 +32,6 @@ interface PaletteGroup { isLoading?: boolean; } -const shortcutChipClass = - "rounded-sm border border-border border-b-2 bg-muted px-[3px] font-mono text-muted-foreground"; - export function CommandPalette({ open, setOpen }: { open: boolean; setOpen: (open: boolean) => void }) { if (!open) return null; return ; @@ -305,13 +303,13 @@ function CommandPaletteDialog({ setOpen }: { setOpen: (open: boolean) => void })
- ↵ 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/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 51a248daafa..6bd7f4e0dbb 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -236,6 +236,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()); @@ -369,6 +370,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 6bda9b4fccb..fc0da6174c4 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -199,6 +199,7 @@ export enum Providers { XINFERENCE = "Xinference", ZAI = "Z.AI (Zhipu AI)", BIGMODEL = "BigModel (Zhipu AI)", + ZeroGPU = "ZeroGPU", } export const provider_map: Record = { @@ -328,6 +329,7 @@ export const provider_map: Record = { XINFERENCE: "xinference", ZAI: "zai", BIGMODEL: "bigmodel", + ZeroGPU: "zerogpu", }; const standaloneSubproviderSlugs = new Set(["bedrock_mantle"]); @@ -502,6 +504,7 @@ const providerPlaceholderMap: Partial> = { [Providers.WATSONX]: "watsonx/ibm/granite-3-3-8b-instruct", [Providers.ZAI]: "zai/glm-4.5", [Providers.BIGMODEL]: "bigmodel/glm-5.3", + [Providers.ZeroGPU]: "zerogpu/gpt-oss-120b", }; export const getPlaceholder = (selectedProvider: string): string => { diff --git a/ui/litellm-dashboard/src/components/shared/ShortcutHints.tsx b/ui/litellm-dashboard/src/components/shared/ShortcutHints.tsx index 6ed5b374cbe..934ba5aa9e0 100644 --- a/ui/litellm-dashboard/src/components/shared/ShortcutHints.tsx +++ b/ui/litellm-dashboard/src/components/shared/ShortcutHints.tsx @@ -1,8 +1,9 @@ "use client"; -import { Fragment, type ReactNode } from "react"; +import { Fragment } from "react"; import { useHotkeysContext } from "react-hotkeys-hook"; +import { Kbd } from "@/components/ui/kbd"; import { cn } from "@/lib/cva.config"; type BoundShortcut = ReturnType["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..70406105193 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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": [ @@ -59856,6 +59857,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 +79055,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 +86392,8 @@ export interface operations { purpose?: string | null; limit?: number | null; after?: string | null; + after_id?: string | null; + before_id?: string | null; }; header?: never; path: {